push

.gitignore

@@ -0,0 +1,63 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual environments
+venv/
+env/
+ENV/
+.venv
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# SMPL models (large files, download separately)
+models/smpl/
+*.pkl
+*.npz
+
+# Trained models
+*.pth
+*.pt
+*.ckpt
+
+# Generated avatars
+*.png
+*.jpg
+*.jpeg
+avatars/
+
+# Logs
+*.log
+logs/
+
+# Environment variables
+.env
+.env.local
+
+# OS
+.DS_Store
+Thumbs.db
+

Dockerfile

@@ -0,0 +1,50 @@
+# Base Image
+FROM python:3.10-slim
+
+# Build argument for Hugging Face token
+ARG HF_TOKEN
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    HF_TOKEN=${HF_TOKEN}
+
+WORKDIR /code
+
+# System Dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    git \
+    curl \
+    libopenblas-dev \
+    libomp-dev \
+    libosmesa6-dev \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements and install Python dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Hugging Face + model tools
+RUN pip install --no-cache-dir huggingface-hub sentencepiece accelerate
+
+# Hugging Face cache environment
+ENV HF_HOME=/models/huggingface \
+    HUGGINGFACE_HUB_CACHE=/models/huggingface \
+    HF_HUB_CACHE=/models/huggingface
+
+# Created cache dir and set permissions
+RUN mkdir -p /models/huggingface && chmod -R 777 /models/huggingface
+
+# Create SMPL models directory
+RUN mkdir -p /code/models/smpl && chmod -R 777 /code/models
+
+# Copy project files
+COPY . .
+
+EXPOSE 7860
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860", "--workers", "1"]
+

app.py

@@ -0,0 +1,93 @@
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel, Field
+from typing import Optional
+import io
+import numpy as np
+from PIL import Image
+
+from measurement_processor import process_measurements
+from beta_regressor import predict_betas
+from smpl_generator import generate_mesh
+from renderer import render_avatar
+
+app = FastAPI(
+    title="Avatar Generation Service",
+    description="Generate 2D avatar images from body measurements using SMPL"
+)
+
+
+class MeasurementRequest(BaseModel):
+    height: float = Field(..., gt=0, description="Height in cm")
+    weight: float = Field(..., gt=0, description="Weight in kg")
+    chest: float = Field(..., gt=0, description="Chest measurement in cm")
+    waist: float = Field(..., gt=0, description="Waist measurement in cm")
+    hips: float = Field(..., gt=0, description="Hips measurement in cm")
+    shoulder_width: Optional[float] = Field(None, gt=0, description="Shoulder width in cm")
+    arm_length: Optional[float] = Field(None, gt=0, description="Arm length in cm")
+    leg_length: Optional[float] = Field(None, gt=0, description="Leg length in cm")
+    inseam: Optional[float] = Field(None, gt=0, description="Inseam in cm")
+    
+    class Config:
+        json_schema_extra = {
+            "example": {
+                "height": 178,
+                "weight": 74,
+                "chest": 96,
+                "waist": 82,
+                "hips": 94,
+                "shoulder_width": 47,
+                "arm_length": 60,
+                "leg_length": 98,
+                "inseam": 81
+            }
+        }
+
+
+@app.get("/")
+async def root():
+    return {
+        "service": "Avatar Generation Service",
+        "endpoints": {
+            "/generate-avatar": "POST - Generate avatar from measurements",
+            "/health": "GET - Health check"
+        }
+    }
+
+
+@app.get("/health")
+async def health():
+    return {"status": "healthy"}
+
+
+@app.post("/generate-avatar")
+async def generate_avatar(measurements: MeasurementRequest):
+    try:
+        measurements_dict = measurements.model_dump(exclude_none=True)
+        normalized = process_measurements(measurements_dict)
+        betas = predict_betas(normalized)
+        vertices, faces = generate_mesh(betas)
+        img_np = render_avatar(vertices, faces)
+        
+        if img_np.dtype != np.uint8:
+            img_np = (img_np * 255).astype(np.uint8) if img_np.max() <= 1.0 else img_np.astype(np.uint8)
+        
+        img = Image.fromarray(img_np, mode='RGB')
+        buf = io.BytesIO()
+        img.save(buf, format="PNG")
+        buf.seek(0)
+        
+        return StreamingResponse(buf, media_type="image/png")
+    
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    except FileNotFoundError as e:
+        raise HTTPException(status_code=500, detail=f"SMPL model not found: {str(e)}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error generating avatar: {str(e)}")
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)
+

beta_regressor.py

@@ -0,0 +1,80 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from pathlib import Path
+from typing import Optional
+
+
+class BetaRegressor(nn.Module):
+    def __init__(self, input_dim: int = 9, output_dim: int = 10, hidden_dims: list = [64, 32]):
+        super(BetaRegressor, self).__init__()
+        
+        layers = []
+        prev_dim = input_dim
+        
+        for hidden_dim in hidden_dims:
+            layers.append(nn.Linear(prev_dim, hidden_dim))
+            layers.append(nn.ReLU())
+            layers.append(nn.Dropout(0.1))
+            prev_dim = hidden_dim
+        
+        layers.append(nn.Linear(prev_dim, output_dim))
+        layers.append(nn.Tanh())
+        
+        self.network = nn.Sequential(*layers)
+        self._initialize_weights()
+    
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight, gain=0.1)
+                nn.init.zeros_(m.bias)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.network(x)
+
+
+class MeasurementToBetaPredictor:
+    def __init__(self, model_path: Optional[str] = None, device: str = "cpu"):
+        self.device = torch.device(device)
+        self.model = BetaRegressor().to(self.device)
+        self.model.eval()
+        
+        if model_path and Path(model_path).exists():
+            self.load_model(model_path)
+        else:
+            print("Warning: Using untrained model. Results may not be optimal.")
+            print("Consider training the model or loading pretrained weights.")
+    
+    def load_model(self, model_path: str):
+        checkpoint = torch.load(model_path, map_location=self.device)
+        if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
+            self.model.load_state_dict(checkpoint['model_state_dict'])
+        else:
+            self.model.load_state_dict(checkpoint)
+        print(f"Loaded model from {model_path}")
+    
+    def predict(self, normalized_measurements: np.ndarray) -> np.ndarray:
+        with torch.no_grad():
+            measurements_tensor = torch.FloatTensor(normalized_measurements).unsqueeze(0).to(self.device)
+            betas_tensor = self.model(measurements_tensor)
+            betas = betas_tensor.squeeze(0).cpu().numpy()
+            betas = betas * 2.0
+        
+        return betas
+
+
+_predictor_instance = None
+
+
+def get_predictor(model_path: Optional[str] = None, device: str = "cpu") -> MeasurementToBetaPredictor:
+    global _predictor_instance
+    if _predictor_instance is None:
+        _predictor_instance = MeasurementToBetaPredictor(model_path=model_path, device=device)
+    return _predictor_instance
+
+
+def predict_betas(normalized_measurements: np.ndarray, model_path: Optional[str] = None) -> np.ndarray:
+    predictor = get_predictor(model_path=model_path)
+    return predictor.predict(normalized_measurements)
+

measurement_processor.py

@@ -0,0 +1,88 @@
+from typing import Dict, Optional
+import numpy as np
+
+
+class MeasurementProcessor:
+    NORMALIZATION_FACTORS = {
+        "height": 200.0,
+        "weight": 100.0,
+        "chest": 150.0,
+        "waist": 150.0,
+        "hips": 150.0,
+        "shoulder_width": 60.0,
+        "arm_length": 80.0,
+        "leg_length": 120.0,
+        "inseam": 100.0,
+    }
+    
+    REQUIRED_MEASUREMENTS = [
+        "height",
+        "weight",
+        "chest",
+        "waist",
+        "hips",
+    ]
+    
+    OPTIONAL_MEASUREMENTS = {
+        "shoulder_width": 40.0,
+        "arm_length": 60.0,
+        "leg_length": 90.0,
+        "inseam": 75.0,
+    }
+    
+    @classmethod
+    def validate_measurements(cls, measurements: Dict) -> Dict[str, str]:
+        errors = []
+        
+        for field in cls.REQUIRED_MEASUREMENTS:
+            if field not in measurements:
+                errors.append(f"Missing required measurement: {field}")
+            elif not isinstance(measurements[field], (int, float)):
+                errors.append(f"Invalid type for {field}: must be number")
+            elif measurements[field] <= 0:
+                errors.append(f"Invalid value for {field}: must be positive")
+        
+        for field, default in cls.OPTIONAL_MEASUREMENTS.items():
+            if field in measurements:
+                if not isinstance(measurements[field], (int, float)):
+                    errors.append(f"Invalid type for {field}: must be number")
+                elif measurements[field] <= 0:
+                    errors.append(f"Invalid value for {field}: must be positive")
+        
+        return {
+            "valid": len(errors) == 0,
+            "errors": errors
+        }
+    
+    @classmethod
+    def normalize_measurements(cls, measurements: Dict) -> np.ndarray:
+        processed = measurements.copy()
+        for field, default in cls.OPTIONAL_MEASUREMENTS.items():
+            if field not in processed:
+                processed[field] = default
+        
+        normalized = []
+        measurement_order = [
+            "height", "weight", "chest", "waist", "hips",
+            "shoulder_width", "arm_length", "leg_length", "inseam"
+        ]
+        
+        for field in measurement_order:
+            value = processed[field]
+            factor = cls.NORMALIZATION_FACTORS[field]
+            normalized.append(value / factor)
+        
+        return np.array(normalized, dtype=np.float32)
+    
+    @classmethod
+    def process(cls, measurements: Dict) -> np.ndarray:
+        validation = cls.validate_measurements(measurements)
+        if not validation["valid"]:
+            raise ValueError(f"Invalid measurements: {', '.join(validation['errors'])}")
+        
+        return cls.normalize_measurements(measurements)
+
+
+def process_measurements(measurements: Dict) -> np.ndarray:
+    return MeasurementProcessor.process(measurements)
+

renderer.py

@@ -0,0 +1,89 @@
+import numpy as np
+import pyrender
+import trimesh
+from typing import Tuple, Optional
+
+
+class AvatarRenderer:
+    def __init__(
+        self,
+        image_size: int = 512,
+        camera_type: str = "orthographic",
+        light_intensity: float = 2.0
+    ):
+        self.image_size = image_size
+        self.camera_type = camera_type
+        self.light_intensity = light_intensity
+        self.renderer = pyrender.OffscreenRenderer(image_size, image_size)
+    
+    def render(
+        self,
+        vertices: np.ndarray,
+        faces: np.ndarray,
+        camera_pose: Optional[np.ndarray] = None,
+        light_pose: Optional[np.ndarray] = None
+    ) -> np.ndarray:
+        mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
+        mesh.vertices -= mesh.vertices.mean(axis=0)
+        scale = 1.0 / (mesh.vertices.max(axis=0) - mesh.vertices.min(axis=0)).max()
+        mesh.vertices *= scale * 0.8
+        
+        pyrender_mesh = pyrender.Mesh.from_trimesh(mesh)
+        scene = pyrender.Scene(bg_color=[1.0, 1.0, 1.0, 1.0])
+        scene.add(pyrender_mesh)
+        
+        if self.camera_type == "orthographic":
+            camera = pyrender.OrthographicCamera(xmag=1.0, ymag=1.0)
+        else:
+            camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0, aspectRatio=1.0)
+        
+        if camera_pose is None:
+            camera_pose = np.eye(4)
+            camera_pose[:3, 3] = [0, 0, 2.5]
+        scene.add(camera, pose=camera_pose)
+        
+        if light_pose is None:
+            light_pose = np.eye(4)
+            light_pose[:3, 3] = [0, 0, 2.5]
+        
+        light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=self.light_intensity)
+        scene.add(light, pose=light_pose)
+        
+        ambient_light = pyrender.SpotLight(
+            color=[1.0, 1.0, 1.0],
+            intensity=0.5,
+            innerConeAngle=np.pi / 16.0,
+            outerConeAngle=np.pi / 6.0
+        )
+        scene.add(ambient_light, pose=light_pose)
+        
+        color, depth = self.renderer.render(scene)
+        
+        if color.shape[2] == 4:
+            color = color[:, :, :3]
+        
+        return color
+    
+    def __del__(self):
+        if hasattr(self, 'renderer'):
+            self.renderer.delete()
+
+
+_renderer_instance = None
+
+
+def get_renderer(image_size: int = 512, camera_type: str = "orthographic") -> AvatarRenderer:
+    global _renderer_instance
+    if _renderer_instance is None:
+        _renderer_instance = AvatarRenderer(image_size=image_size, camera_type=camera_type)
+    return _renderer_instance
+
+
+def render_avatar(
+    vertices: np.ndarray,
+    faces: np.ndarray,
+    image_size: int = 512
+) -> np.ndarray:
+    renderer = get_renderer(image_size=image_size)
+    return renderer.render(vertices, faces)
+

requirements.txt

@@ -0,0 +1,11 @@
+fastapi
+uvicorn[standard]
+pydantic
+numpy
+torch
+smplx
+pyrender
+trimesh
+Pillow
+scikit-learn
+

smpl_generator.py

@@ -0,0 +1,105 @@
+import torch
+import numpy as np
+from pathlib import Path
+from typing import Tuple, Optional
+import smplx
+
+
+class SMPLGenerator:
+    def __init__(self, model_path: str = "models/smpl", gender: str = "neutral", device: str = "cpu"):
+        self.device = torch.device(device)
+        self.gender = gender
+        self.model_path = Path(model_path)
+        
+        if not self.model_path.exists():
+            raise FileNotFoundError(
+                f"SMPL model not found at {model_path}. "
+                f"Please download SMPL models from https://smpl.is.tue.mpg.de/"
+            )
+        
+        try:
+            self.smpl_model = smplx.create(
+                str(self.model_path),
+                model_type='smpl',
+                gender=gender,
+                batch_size=1,
+                ext='npz'
+            ).to(self.device)
+        except Exception as e:
+            try:
+                self.smpl_model = smplx.create(
+                    str(self.model_path),
+                    model_type='smpl',
+                    gender=gender,
+                    batch_size=1,
+                    ext='pkl'
+                ).to(self.device)
+            except Exception as e2:
+                raise RuntimeError(
+                    f"Failed to load SMPL model: {e}. "
+                    f"Please ensure SMPL model files are in {model_path}"
+                )
+    
+    def generate_mesh(
+        self,
+        betas: np.ndarray,
+        body_pose: Optional[np.ndarray] = None,
+        global_orient: Optional[np.ndarray] = None,
+        transl: Optional[np.ndarray] = None
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        if betas.ndim == 1:
+            betas = betas.unsqueeze(0) if isinstance(betas, torch.Tensor) else betas[np.newaxis, :]
+        
+        if isinstance(betas, np.ndarray):
+            betas = torch.FloatTensor(betas).to(self.device)
+        
+        batch_size = betas.shape[0]
+        
+        if body_pose is None:
+            body_pose = torch.zeros([batch_size, 69], device=self.device)
+        elif isinstance(body_pose, np.ndarray):
+            body_pose = torch.FloatTensor(body_pose).to(self.device)
+        
+        if global_orient is None:
+            global_orient = torch.zeros([batch_size, 3], device=self.device)
+        elif isinstance(global_orient, np.ndarray):
+            global_orient = torch.FloatTensor(global_orient).to(self.device)
+        
+        if transl is None:
+            transl = torch.zeros([batch_size, 3], device=self.device)
+        elif isinstance(transl, np.ndarray):
+            transl = torch.FloatTensor(transl).to(self.device)
+        
+        with torch.no_grad():
+            output = self.smpl_model(
+                betas=betas,
+                body_pose=body_pose,
+                global_orient=global_orient,
+                transl=transl
+            )
+        
+        vertices = output.vertices[0].detach().cpu().numpy()
+        faces = self.smpl_model.faces
+        
+        return vertices, faces
+
+
+_generator_instance = None
+
+
+def get_generator(model_path: str = "models/smpl", gender: str = "neutral", device: str = "cpu") -> SMPLGenerator:
+    global _generator_instance
+    if _generator_instance is None:
+        _generator_instance = SMPLGenerator(model_path=model_path, gender=gender, device=device)
+    return _generator_instance
+
+
+def generate_mesh(
+    betas: np.ndarray,
+    model_path: str = "models/smpl",
+    gender: str = "neutral",
+    device: str = "cpu"
+) -> Tuple[np.ndarray, np.ndarray]:
+    generator = get_generator(model_path=model_path, gender=gender, device=device)
+    return generator.generate_mesh(betas)
+