Spaces:
Running
Running
| 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) | |
| async def root(): | |
| """ | |
| Root endpoint serving a simple HTML page with API documentation. | |
| """ | |
| return landing_html | |
| # API endpoints | |
| 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)) | |
| 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)) | |
| 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"] | |
| ) | |
| 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() | |
| ] | |
| 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) |