Spaces:

jakereardon
/

mlse-player-3d

Running

File size: 11,949 Bytes

import os
import io
import base64
import uuid
import asyncio
from typing import Optional, Dict, Any, List
from fastapi import FastAPI, HTTPException, Body, BackgroundTasks, File, UploadFile, Form, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel
import uvicorn
import logging

from app.utils import ensure_directories

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(
    title="MLSE Player 3D Generator",
    description="API for generating 3D human body models from player images using SAM 3D Body",
    version="0.1.0"
)

# Add CORS middleware for frontend integration
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Update this with specific origins in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Create required directories
ensure_directories()

# Mount static files directory if it exists
if os.path.exists("outputs"):
    app.mount("/outputs", StaticFiles(directory="outputs"), name="outputs")

# In-memory job storage (replace with database in production)
jobs = {}

# Define HTML content for the root page
landing_html = """
<!DOCTYPE html>
<html>
<head>
    <title>MLSE Player 3D Generator</title>
    <style>
        body {
            font-family: Arial, sans-serif;
            line-height: 1.6;
            max-width: 800px;
            margin: 0 auto;
            padding: 20px;
        }
        h1 {
            color: #2563eb;
            border-bottom: 2px solid #e5e7eb;
            padding-bottom: 10px;
        }
        h2 {
            color: #1e40af;
            margin-top: 20px;
        }
        code {
            background-color: #f3f4f6;
            padding: 2px 5px;
            border-radius: 3px;
            font-family: monospace;
        }
        pre {
            background-color: #f3f4f6;
            padding: 15px;
            border-radius: 5px;
            overflow-x: auto;
        }
        .endpoint {
            margin-bottom: 20px;
            padding: 10px;
            border-left: 3px solid #2563eb;
            background-color: #f9fafb;
        }
        .method {
            font-weight: bold;
            color: #2563eb;
        }
    </style>
</head>
<body>
    <h1>MLSE Player 3D Generator</h1>
    <p>A 3D player model generator that uses AI to convert images of athletes into detailed 3D models.</p>

    <h2>API Endpoints</h2>

    <div class="endpoint">
        <p><span class="method">POST</span> /api/upload</p>
        <p>Upload an image file to generate a 3D model</p>
        <p><strong>Form Data:</strong></p>
        <ul>
            <li><code>file</code>: Image file (JPEG/PNG)</li>
            <li><code>player_name</code>: Name for the model (default: "player")</li>
            <li><code>use_keypoints</code>: Whether to use keypoint detection (default: true)</li>
            <li><code>use_mask</code>: Whether to use segmentation masks (default: true)</li>
        </ul>
    </div>

    <div class="endpoint">
        <p><span class="method">POST</span> /api/process</p>
        <p>Process a base64-encoded image</p>
        <p><strong>JSON Body:</strong></p>
        <pre>
{
  "image_data": "base64_encoded_image_data",
  "player_name": "player_name",
  "options": {
    "use_keypoints": true,
    "use_mask": true
  }
}
        </pre>
    </div>

    <div class="endpoint">
        <p><span class="method">POST</span> /api/status</p>
        <p>Check the status of a processing job</p>
        <p><strong>JSON Body:</strong></p>
        <pre>
{
  "job_id": "job_id_from_upload_response"
}
        </pre>
    </div>

    <div class="endpoint">
        <p><span class="method">GET</span> /api/jobs</p>
        <p>List all processing jobs</p>
    </div>

    <div class="endpoint">
        <p><span class="method">GET</span> /api/model/{job_id}</p>
        <p>Get the 3D model file for a completed job</p>
    </div>

    <p>Note: This is a demo version using simplified mock processing. For the full version with SAM 3D Body integration, additional setup is required.</p>

    <footer>
        <p>MLSE Player 3D Generator | Powered by Hugging Face Spaces</p>
    </footer>
</body>
</html>
"""

# Request models
class ImageProcessRequest(BaseModel):
    image_data: str  # Base64 encoded image
    player_name: str = "player"  # Name for the generated model
    options: Dict[str, Any] = {
        "use_keypoints": True,
        "use_mask": True
    }

    class Config:
        protected_namespaces = ()  # Fix for "model_" namespace warning

class JobStatusRequest(BaseModel):
    job_id: str

# Response models
class JobResponse(BaseModel):
    job_id: str
    status: str = "queued"  # queued, processing, completed, failed

class JobStatusResponse(BaseModel):
    job_id: str
    status: str
    progress: float = 0
    model_url: Optional[str] = None
    preview_url: Optional[str] = None
    error: Optional[str] = None

    class Config:
        protected_namespaces = ()  # Fix for "model_" namespace warning

# Initialize the model on startup (using a context manager instead of on_event)
@app.get("/", response_class=HTMLResponse)
async def root():
    """
    Root endpoint serving a simple HTML page with API documentation.
    """
    return landing_html

# API endpoints
@app.post("/api/process", response_model=JobResponse)
async def process_image_endpoint(request: ImageProcessRequest, background_tasks: BackgroundTasks):
    """
    Process an image to generate a 3D model using SAM 3D Body.
    Accepts a base64-encoded image and returns a job ID for tracking progress.
    """
    try:
        # Generate a unique job ID
        job_id = str(uuid.uuid4())

        # Store job in memory
        jobs[job_id] = {
            "status": "queued",
            "progress": 0,
            "model_url": None,
            "preview_url": None,
            "error": None
        }

        # Process in background
        background_tasks.add_task(
            process_image_background,
            job_id,
            request.image_data,
            request.player_name,
            request.options
        )

        return JobResponse(job_id=job_id)

    except Exception as e:
        logger.error(f"Error processing image: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/upload", response_model=JobResponse)
async def upload_image_endpoint(
    file: UploadFile = File(...),
    player_name: str = Form("player"),
    use_keypoints: bool = Form(True),
    use_mask: bool = Form(True),
    background_tasks: BackgroundTasks = None
):
    """
    Process an uploaded image to generate a 3D model.
    This endpoint accepts multipart/form-data for easier frontend integration.
    """
    try:
        # Generate a unique job ID
        job_id = str(uuid.uuid4())

        # Read the image file
        image_bytes = await file.read()

        # Convert to base64 for consistency with the other endpoint
        image_data = base64.b64encode(image_bytes).decode('utf-8')

        # Store job in memory
        jobs[job_id] = {
            "status": "queued",
            "progress": 0,
            "model_url": None,
            "preview_url": None,
            "error": None
        }

        # Process in background
        options = {
            "use_keypoints": use_keypoints,
            "use_mask": use_mask
        }
        background_tasks.add_task(
            process_image_background,
            job_id,
            image_data,
            player_name,
            options
        )

        return JobResponse(job_id=job_id)

    except Exception as e:
        logger.error(f"Error uploading image: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/status", response_model=JobStatusResponse)
async def check_status_endpoint(request: JobStatusRequest):
    """
    Check the status of a processing job by job ID.
    """
    job_id = request.job_id
    if job_id not in jobs:
        raise HTTPException(status_code=404, detail=f"Job {job_id} not found")

    job_info = jobs[job_id]
    return JobStatusResponse(
        job_id=job_id,
        status=job_info["status"],
        progress=job_info["progress"],
        model_url=job_info["model_url"],
        preview_url=job_info["preview_url"],
        error=job_info["error"]
    )

@app.get("/api/jobs", response_model=List[JobStatusResponse])
async def list_jobs_endpoint():
    """
    List all processing jobs and their status.
    """
    return [
        JobStatusResponse(
            job_id=job_id,
            status=job_info["status"],
            progress=job_info["progress"],
            model_url=job_info["model_url"],
            preview_url=job_info["preview_url"],
            error=job_info["error"]
        )
        for job_id, job_info in jobs.items()
    ]

@app.get("/api/model/{job_id}")
async def get_model_endpoint(job_id: str):
    """
    Get the 3D model file for a completed job.
    """
    if job_id not in jobs:
        raise HTTPException(status_code=404, detail=f"Job {job_id} not found")

    job_info = jobs[job_id]
    if job_info["status"] != "completed" or not job_info["model_url"]:
        raise HTTPException(status_code=400, detail="Model not ready or failed")

    # Return the model file
    model_path = job_info["model_url"].replace("/outputs/", "outputs/")
    return FileResponse(model_path)

# Background task for processing images
async def process_image_background(job_id, image_data, player_name, options):
    try:
        # Update job status
        jobs[job_id]["status"] = "processing"
        jobs[job_id]["progress"] = 10

        # Decode base64 image if needed
        if isinstance(image_data, str) and image_data.startswith('data:image'):
            image_data = image_data.split(',')[1]

        if isinstance(image_data, str):
            image_bytes = base64.b64decode(image_data)
        else:
            image_bytes = image_data

        # Save to temporary file
        os.makedirs("temp", exist_ok=True)
        input_path = f"temp/{job_id}_input.jpg"
        with open(input_path, 'wb') as f:
            f.write(image_bytes)

        jobs[job_id]["progress"] = 20

        # Process the image with SAM 3D Body
        from app.sam_3d_service import process_image
        result = await asyncio.to_thread(
            process_image,
            input_path,
            player_name,
            options.get("use_keypoints", True),
            options.get("use_mask", True),
            lambda progress: update_job_progress(job_id, progress)
        )

        # Update job with result
        model_path = result["model_path"]
        preview_path = result["preview_path"]

        jobs[job_id].update({
            "status": "completed",
            "progress": 100,
            "model_url": f"/outputs/{job_id}/{player_name}.glb",
            "preview_url": f"/outputs/{job_id}/{player_name}_preview.jpg"
        })

    except Exception as e:
        logger.error(f"Error processing job {job_id}: {str(e)}")
        jobs[job_id].update({
            "status": "failed",
            "error": str(e)
        })

# No longer needed as we use the real SAM 3D Body implementation now

def update_job_progress(job_id: str, progress: float):
    """Update the progress of a job"""
    if job_id in jobs:
        # Scale progress to 20-90% range (we reserve 0-20% for setup and 90-100% for final steps)
        scaled_progress = 20 + (progress * 70)
        jobs[job_id]["progress"] = min(90, scaled_progress)

# Serve the app with uvicorn if run directly
if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run("app.main:app", host="0.0.0.0", port=port)