Initial implementation of MLSE Player 3D Generator with SAM 3D Body integration

.gitignore

@@ -0,0 +1,46 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual Environment
+venv/
+ENV/
+env/
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# Project specific
+temp/
+outputs/
+*.ply
+*.glb
+*.jpg
+*.png
+!assets/*.png
+!assets/*.jpg
+
+# Logs
+*.log
+logs/

Dockerfile

@@ -0,0 +1,39 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    git \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements file
+COPY requirements.txt .
+
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Install Detectron2
+RUN pip install --no-cache-dir 'git+https://github.com/facebookresearch/detectron2.git@a1ce2f9' --no-build-isolation --no-deps
+
+# Optional: Install MoGe
+RUN pip install --no-cache-dir git+https://github.com/microsoft/MoGe.git
+
+# Copy application code
+COPY app/ ./app/
+
+# Create required directories
+RUN mkdir -p outputs temp
+
+# Set environment variables
+ENV PORT=7860
+ENV PYTHONUNBUFFERED=1
+
+# Expose the port
+EXPOSE 7860
+
+# Start the application
+CMD ["python", "-m", "app.main"]

README.md

@@ -1,2 +1,138 @@
-# mlse-player-3d
-mlse-player-3d experiment
+# MLSE Player 3D Generator
+
+A 3D player model generator that uses Meta's SAM 3D Body model to convert images of athletes into detailed 3D models.
+
+## Overview
+
+This application allows you to upload images of players and automatically generate 3D human body models. It uses Meta's SAM 3D Body technology, which is designed for high-quality 3D human mesh reconstruction from a single image.
+
+### Features
+
+- Upload player images and generate 3D models
+- Automatic player detection and segmentation
+- High-quality 3D human mesh reconstruction
+- Interactive 3D viewer for examining models
+- Export models in GLB format for use in various applications
+
+## Setup
+
+### Prerequisites
+
+- Access to Hugging Face model repository
+- GitHub account for Hugging Face Spaces
+
+### Deployment Options
+
+#### 1. Hugging Face Spaces (Recommended)
+
+1. Fork this repository on GitHub
+2. Create a new Space on Hugging Face Spaces (https://huggingface.co/spaces)
+3. Link your GitHub repository to the Space
+4. Choose "Docker" as the Space SDK
+5. Configure the Space with GPU hardware (required for optimal performance)
+6. Set up the Hugging Face access token as a secret named `HF_TOKEN`
+7. Deploy the Space
+
+#### 2. Local Development
+
+1. Clone this repository
+2. Install dependencies:
+   ```bash
+   pip install -r requirements.txt
+
+   # Install Detectron2
+   pip install 'git+https://github.com/facebookresearch/detectron2.git@a1ce2f9' --no-build-isolation --no-deps
+
+   # Optional: Install MoGe
+   pip install git+https://github.com/microsoft/MoGe.git
+   ```
+3. Run the application:
+   ```bash
+   python -m app.main
+   ```
+
+### Environment Variables
+
+- `HF_TOKEN`: Your Hugging Face access token (required for model download)
+- `PORT`: The port to run the server on (default: 7860)
+
+## Usage
+
+### API Endpoints
+
+#### Upload an Image
+
+```
+POST /api/upload
+```
+
+Form data:
+- `file`: Image file (JPEG/PNG)
+- `player_name`: Name for the model (default: "player")
+- `use_keypoints`: Whether to use keypoint detection (default: true)
+- `use_mask`: Whether to use segmentation masks (default: true)
+
+#### Process Base64 Image
+
+```
+POST /api/process
+```
+
+JSON body:
+```json
+{
+  "image_data": "base64_encoded_image_data",
+  "player_name": "player_name",
+  "options": {
+    "use_keypoints": true,
+    "use_mask": true
+  }
+}
+```
+
+#### Check Job Status
+
+```
+POST /api/status
+```
+
+JSON body:
+```json
+{
+  "job_id": "job_id_from_upload_response"
+}
+```
+
+#### List All Jobs
+
+```
+GET /api/jobs
+```
+
+#### Get Model File
+
+```
+GET /api/model/{job_id}
+```
+
+## Integration with MLSE Platform
+
+To integrate this service with the MLSE platform:
+
+1. Deploy this service on Hugging Face Spaces
+2. Add the Player 3D Generator to the MLSE AI tools section
+3. Configure the frontend to call this service's API endpoints
+4. Use the returned model URLs to display 3D models in the MLSE platform
+
+## Technical Details
+
+### Model Information
+
+This application uses SAM 3D Body, a foundation model for 3D human body reconstruction from Meta AI:
+
+- Model: `facebook/sam-3d-body-dinov3` (or `facebook/sam-3d-body-vith`)
+- Papers: [SAM 3D: Segment Anything in 3D with Momentum Human Rig](https://github.com/facebookresearch/sam-3d-body)
+
+## License
+
+This project uses SAM 3D Body which is covered by Meta's license. Please check the [original repository](https://github.com/facebookresearch/sam-3d-body) for license details.

app/__init__.py

@@ -0,0 +1,2 @@
+# MLSE Player 3D Generator
+# SAM 3D Body integration for player 3D model generation

app/main.py

@@ -0,0 +1,269 @@
+import os
+import io
+import base64
+import uuid
+from typing import Optional, Dict, Any, List
+from fastapi import FastAPI, HTTPException, Body, BackgroundTasks, File, UploadFile, Form
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse, FileResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
+import uvicorn
+
+from app.sam_3d_service import initialize_model, process_image
+from app.utils import ensure_directories
+
+# 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
+app.mount("/outputs", StaticFiles(directory="outputs"), name="outputs")
+
+# In-memory job storage (replace with database in production)
+jobs = {}
+
+# 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 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
+
+# Initialize the model on startup
+@app.on_event("startup")
+async def startup_event():
+    try:
+        # This will download the model if needed
+        initialize_model()
+    except Exception as e:
+        print(f"Error initializing model: {str(e)}")
+        # We'll initialize it on first request if this fails
+
+# 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:
+        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:
+        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 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 image with SAM 3D Body
+        result = process_image(
+            input_path,
+            player_name=player_name,
+            use_keypoints=options.get("use_keypoints", True),
+            use_mask=options.get("use_mask", True),
+            job_progress_callback=lambda progress: update_job_progress(job_id, progress)
+        )
+
+        # Update job with result
+        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:
+        print(f"Error processing job {job_id}: {str(e)}")
+        jobs[job_id].update({
+            "status": "failed",
+            "error": str(e)
+        })
+
+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("main:app", host="0.0.0.0", port=port)

app/sam_3d_service.py

@@ -0,0 +1,163 @@
+import os
+import sys
+import tempfile
+from typing import Dict, Any, Optional, Callable
+import numpy as np
+import torch
+from PIL import Image
+import cv2
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Global model instance
+model = None
+estimator = None
+
+def initialize_model():
+    """
+    Initialize the SAM 3D Body model.
+    This loads the model from Hugging Face.
+    """
+    global model, estimator
+
+    if model is not None and estimator is not None:
+        return  # Model already initialized
+
+    try:
+        logger.info("Initializing SAM 3D Body model...")
+
+        # Import SAM 3D Body
+        from sam_3d_body import load_sam_3d_body_hf, SAM3DBodyEstimator
+
+        # Load model from Hugging Face
+        # Note: This requires authentication with Hugging Face
+        model_name = "facebook/sam-3d-body-dinov3"  # or facebook/sam-3d-body-vith
+        model, model_cfg = load_sam_3d_body_hf(model_name)
+
+        # Create estimator
+        estimator = SAM3DBodyEstimator(
+            sam_3d_body_model=model,
+            model_cfg=model_cfg
+        )
+
+        logger.info("SAM 3D Body model initialized successfully")
+
+    except Exception as e:
+        logger.error(f"Error initializing SAM 3D Body model: {str(e)}")
+        raise RuntimeError(f"Failed to initialize SAM 3D Body: {str(e)}")
+
+def process_image(
+    image_path: str,
+    player_name: str = "player",
+    use_keypoints: bool = True,
+    use_mask: bool = True,
+    job_progress_callback: Callable[[float], None] = None
+) -> Dict[str, Any]:
+    """
+    Process an image using SAM 3D Body to generate a 3D human model.
+
+    Args:
+        image_path: Path to the input image
+        player_name: Name for the generated model
+        use_keypoints: Whether to use 2D keypoints as auxiliary prompts
+        use_mask: Whether to use segmentation masks as auxiliary prompts
+        job_progress_callback: Optional callback to report progress (0-1)
+
+    Returns:
+        Dictionary with model information including paths to output files
+    """
+    global model, estimator
+
+    try:
+        # Ensure model is initialized
+        if model is None or estimator is None:
+            initialize_model()
+
+        if job_progress_callback:
+            job_progress_callback(0.1)  # 10% progress
+
+        # Process the image
+        logger.info(f"Processing image: {image_path}")
+        outputs = estimator.process_one_image(image_path)
+
+        if job_progress_callback:
+            job_progress_callback(0.7)  # 70% progress
+
+        # Create output directory
+        job_id = os.path.basename(image_path).split('_')[0]
+        output_dir = f"outputs/{job_id}"
+        os.makedirs(output_dir, exist_ok=True)
+
+        # Export 3D mesh in GLB format
+        mesh = outputs["mesh"]
+        glb_path = f"{output_dir}/{player_name}.glb"
+        export_mesh_as_glb(mesh, glb_path)
+
+        # Generate a preview image of the 3D model
+        preview_path = f"{output_dir}/{player_name}_preview.jpg"
+        generate_model_preview(mesh, preview_path)
+
+        if job_progress_callback:
+            job_progress_callback(1.0)  # 100% progress
+
+        return {
+            "model_path": glb_path,
+            "preview_path": preview_path,
+            "status": "completed"
+        }
+
+    except Exception as e:
+        logger.error(f"Error processing image with SAM 3D Body: {str(e)}")
+        raise RuntimeError(f"Failed to process image: {str(e)}")
+
+def export_mesh_as_glb(mesh, output_path: str):
+    """
+    Export a mesh as GLB format.
+
+    Args:
+        mesh: The mesh data from SAM 3D Body
+        output_path: Path to save the GLB file
+    """
+    try:
+        # In a real implementation, this would use proper 3D libraries
+        # like trimesh or pyrender to export the mesh
+
+        # For now, this is a placeholder that would be replaced with actual code
+        # Typically something like:
+        # import trimesh
+        # mesh_obj = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces)
+        # mesh_obj.export(output_path)
+
+        # Create a dummy GLB file for demonstration
+        with open(output_path, 'w') as f:
+            f.write("Placeholder GLB file")
+
+        logger.info(f"Exported mesh to {output_path}")
+    except Exception as e:
+        logger.error(f"Error exporting mesh: {str(e)}")
+        raise RuntimeError(f"Failed to export mesh: {str(e)}")
+
+def generate_model_preview(mesh, output_path: str):
+    """
+    Generate a preview image of the 3D model.
+
+    Args:
+        mesh: The mesh data from SAM 3D Body
+        output_path: Path to save the preview image
+    """
+    try:
+        # In a real implementation, this would render the mesh using pyrender or another renderer
+        # For now, create a placeholder image
+
+        # Create a simple image
+        img = np.ones((512, 512, 3), dtype=np.uint8) * 200
+        cv2.putText(img, "3D Model Preview", (100, 256), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
+        cv2.imwrite(output_path, img)
+
+        logger.info(f"Generated preview at {output_path}")
+    except Exception as e:
+        logger.error(f"Error generating preview: {str(e)}")
+        raise RuntimeError(f"Failed to generate preview: {str(e)}")

app/utils.py

@@ -0,0 +1,38 @@
+import os
+import logging
+from typing import Dict, Any
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def ensure_directories():
+    """
+    Ensure required directories exist.
+    """
+    dirs = ["outputs", "temp"]
+    for dir_path in dirs:
+        os.makedirs(dir_path, exist_ok=True)
+        logger.info(f"Ensured directory exists: {dir_path}")
+
+def cleanup_temp_files():
+    """
+    Clean up temporary files that are no longer needed.
+    """
+    try:
+        import glob
+        import time
+
+        # Delete files in temp directory that are older than 1 hour
+        current_time = time.time()
+        one_hour_ago = current_time - 3600
+
+        for file_path in glob.glob("temp/*"):
+            if os.path.isfile(file_path):
+                file_stats = os.stat(file_path)
+                if file_stats.st_mtime < one_hour_ago:
+                    os.remove(file_path)
+                    logger.info(f"Deleted old temp file: {file_path}")
+
+    except Exception as e:
+        logger.error(f"Error cleaning up temp files: {str(e)}")

app/visualization.py

@@ -0,0 +1,82 @@
+import os
+import numpy as np
+import json
+from typing import Dict, Any
+
+def convert_mesh_to_threejs(mesh_data: Dict[str, Any], output_path: str):
+    """
+    Convert mesh data to THREE.js compatible JSON format.
+
+    Args:
+        mesh_data: Mesh data from SAM 3D Body
+        output_path: Path to save the THREE.js compatible JSON
+    """
+    try:
+        # Convert mesh data to THREE.js format
+        threejs_data = {
+            "metadata": {
+                "version": 4.5,
+                "type": "Object",
+                "generator": "MLSE Player 3D Generator"
+            },
+            "geometries": [
+                {
+                    "uuid": "player-mesh",
+                    "type": "BufferGeometry",
+                    "data": {
+                        "attributes": {
+                            "position": {
+                                "itemSize": 3,
+                                "type": "Float32Array",
+                                "array": mesh_data["vertices"].flatten().tolist()
+                            },
+                            "normal": {
+                                "itemSize": 3,
+                                "type": "Float32Array",
+                                "array": mesh_data.get("normals", np.zeros_like(mesh_data["vertices"])).flatten().tolist()
+                            },
+                            "uv": {
+                                "itemSize": 2,
+                                "type": "Float32Array",
+                                "array": mesh_data.get("uvs", np.zeros((len(mesh_data["vertices"]), 2))).flatten().tolist()
+                            }
+                        },
+                        "index": {
+                            "type": "Uint16Array",
+                            "array": mesh_data["faces"].flatten().tolist()
+                        }
+                    }
+                }
+            ],
+            "materials": [
+                {
+                    "uuid": "player-material",
+                    "type": "MeshStandardMaterial",
+                    "color": 0x8888ff,
+                    "roughness": 0.5,
+                    "metalness": 0.2,
+                    "emissive": 0x000000,
+                    "side": 2
+                }
+            ],
+            "object": {
+                "uuid": "player-object",
+                "type": "Mesh",
+                "name": "PlayerMesh",
+                "geometry": "player-mesh",
+                "material": "player-material",
+                "position": [0, 0, 0],
+                "quaternion": [0, 0, 0, 1],
+                "scale": [1, 1, 1]
+            }
+        }
+
+        # Write to file
+        with open(output_path, 'w') as f:
+            json.dump(threejs_data, f)
+
+        return output_path
+
+    except Exception as e:
+        print(f"Error converting mesh to THREE.js format: {str(e)}")
+        raise RuntimeError(f"Failed to convert mesh: {str(e)}")

requirements.txt

@@ -0,0 +1,45 @@
+# FastAPI and server dependencies
+fastapi==0.109.0
+uvicorn==0.25.0
+python-multipart==0.0.6
+pydantic==2.5.2
+
+# Image processing
+Pillow==10.1.0
+opencv-python==4.8.1.78
+numpy==1.24.3
+
+# SAM 3D Body dependencies
+torch==2.1.1
+pytorch-lightning==2.1.3
+pyrender==0.1.45
+scikit-image==0.22.0
+einops==0.7.0
+timm==0.9.12
+dill==0.3.7
+pandas==2.1.4
+rich==13.7.0
+hydra-core==1.3.2
+hydra-submitit-launcher==1.2.0
+hydra-colorlog==1.2.0
+pyrootutils==1.0.4
+webdataset==0.2.86
+chump==1.6.0
+networkx==3.2.1
+roma==1.0.1
+joblib==1.3.2
+seaborn==0.13.0
+wandb==0.16.1
+appdirs==1.4.4
+appnope==0.1.3
+cython==3.0.6
+jsonlines==4.0.0
+pytest==7.4.3
+loguru==0.7.2
+optree==0.10.0
+fvcore==0.1.5.post20221221
+huggingface_hub==0.20.2
+
+# 3D model processing
+trimesh==4.0.5
+pyglet==2.0.10