push

app.py

@@ -48,6 +48,7 @@ class MeasurementRequest(BaseModel):
     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")
+    gender: Optional[str] = Field("male", description="Gender: 'male' or 'female'")
     
     class Config:
         json_schema_extra = {
@@ -60,7 +61,8 @@ class MeasurementRequest(BaseModel):
                 "shoulder_width": 47,
                 "arm_length": 60,
                 "leg_length": 98,
-                "inseam": 81
+                "inseam": 81,
+                "gender": "male"
             }
         }
 
@@ -70,7 +72,8 @@ async def root():
     return {
         "service": "Avatar Generation Service",
         "endpoints": {
-            "/generate-avatar": "POST - Generate avatar from measurements",
+            "/generate-avatar": "POST - Generate 2D avatar image (PNG) from measurements",
+            "/generate-avatar-3d": "POST - Generate 3D avatar mesh (OBJ) from measurements",
             "/health": "GET - Health check"
         }
     }
@@ -85,9 +88,14 @@ async def health():
 async def generate_avatar(measurements: MeasurementRequest):
     try:
         measurements_dict = measurements.model_dump(exclude_none=True)
+        gender = measurements_dict.pop("gender", "male")
+        
+        if gender not in ["male", "female", "neutral"]:
+            raise ValueError("Gender must be 'male', 'female', or 'neutral'")
+        
         normalized = process_measurements(measurements_dict)
         betas = predict_betas(normalized)
-        vertices, faces = generate_mesh(betas, model_path=SMPL_MODEL_PATH)
+        vertices, faces = generate_mesh(betas, model_path=SMPL_MODEL_PATH, gender=gender)
         img_np = render_avatar(vertices, faces)
         
         if img_np.dtype != np.uint8:
@@ -113,6 +121,45 @@ async def generate_avatar(measurements: MeasurementRequest):
         raise HTTPException(status_code=500, detail=f"Error generating avatar: {str(e)}")
 
 
+@app.post("/generate-avatar-3d")
+async def generate_avatar_3d(measurements: MeasurementRequest):
+    try:
+        measurements_dict = measurements.model_dump(exclude_none=True)
+        gender = measurements_dict.pop("gender", "male")
+        
+        if gender not in ["male", "female", "neutral"]:
+            raise ValueError("Gender must be 'male', 'female', or 'neutral'")
+        
+        normalized = process_measurements(measurements_dict)
+        betas = predict_betas(normalized)
+        vertices, faces = generate_mesh(betas, model_path=SMPL_MODEL_PATH, gender=gender)
+        
+        import trimesh
+        mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
+        
+        buf = io.BytesIO()
+        mesh.export(file_obj=buf, file_type='obj')
+        buf.seek(0)
+        
+        return StreamingResponse(
+            buf,
+            media_type="model/obj",
+            headers={"Content-Disposition": "attachment; filename=avatar.obj"}
+        )
+    
+    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)}. "
+                   f"Please ensure SMPL model files are in {SMPL_MODEL_PATH}. "
+                   f"Download from https://smpl.is.tue.mpg.de/ or set SMPL_MODEL_PATH environment variable."
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error generating 3D avatar: {str(e)}")
+
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)