smpl_generator.py

import torch
import numpy as np
from pathlib import Path
from typing import Tuple, Optional
import smplx
import os


class SMPLGenerator:
    def __init__(self, model_path: str = "smpl", gender: str = "neutral", device: str = "cpu"):
        self.device = torch.device(device)
        self.gender = gender
        
        model_path_obj = Path(model_path)
        
        if not model_path_obj.exists():
            alt_paths = [Path("smpl"), Path("smpl/smpl")]
            for alt_path in alt_paths:
                if alt_path.exists():
                    model_path_obj = alt_path
                    print(f"Using alternative model path: {model_path_obj}")
                    break
            else:
                model_path_obj.mkdir(parents=True, exist_ok=True)
        
        models_source = Path("smpl/smpl/models")
        if not models_source.exists():
            models_source = model_path_obj / "models"
        
        self.model_path = model_path_obj
        model_path_str = str(self.model_path)
        
        if gender == "neutral":
            gender = "male"
            print("Note: Neutral gender not available, using male model")
        
        if models_source.exists():
            model_files = list(models_source.glob("*.pkl"))
            print(f"Found {len(model_files)} model files in {models_source}: {[f.name for f in model_files]}")
            
            import shutil
            
            expected_smpl_dir = Path("smpl") / "smpl"
            expected_models_dir = expected_smpl_dir / "models"
            expected_models_dir.mkdir(parents=True, exist_ok=True)
            
            for model_file in model_files:
                file_lower = model_file.name.lower()
                target_name = None
                
                if "basicmodel_m" in file_lower or "male" in file_lower:
                    target_name = "SMPL_MALE.pkl"
                elif "basicmodel_f" in file_lower or "female" in file_lower:
                    target_name = "SMPL_FEMALE.pkl"
                elif "neutral" in file_lower:
                    target_name = "SMPL_NEUTRAL.pkl"
                
                if target_name:
                    target_in_models = expected_models_dir / target_name
                    target_in_smpl = expected_smpl_dir / target_name
                    
                    if not target_in_models.exists():
                        shutil.copy2(model_file, target_in_models)
                        print(f"Copied {model_file.name} -> {target_in_models}")
                    
                    if not target_in_smpl.exists():
                        shutil.copy2(model_file, target_in_smpl)
                        print(f"Copied {model_file.name} -> {target_in_smpl}")
                else:
                    target_file = expected_models_dir / model_file.name
                    if not target_file.exists():
                        shutil.copy2(model_file, target_file)
                        print(f"Copied {model_file.name} to {target_file}")
        
        models_dir = model_path_obj / "smpl" / "models"
        if not models_dir.exists():
            models_dir = model_path_obj / "models"
        
        base_path = Path(".").absolute()
        model_paths_to_try = [
            str(base_path),
            ".",
            "smpl",
            str(model_path_obj),
        ]
        
        if models_dir.exists():
            parent_of_smpl = models_dir.parent.parent
            if parent_of_smpl.exists():
                model_paths_to_try.append(str(parent_of_smpl))
        
        model_paths_to_try = list(dict.fromkeys(model_paths_to_try))
        
        last_error = None
        for try_path in model_paths_to_try:
            print(f"Trying model path: {try_path}")
            try:
                self.smpl_model = smplx.create(
                    model_path=try_path,
                    model_type='smpl',
                    gender=gender,
                    batch_size=1,
                    ext='npz'
                ).to(self.device)
                print(f"Successfully loaded model from: {try_path}")
                break
            except Exception as e:
                last_error = e
                try:
                    self.smpl_model = smplx.create(
                        model_path=try_path,
                        model_type='smpl',
                        gender=gender,
                        batch_size=1,
                        ext='pkl'
                    ).to(self.device)
                    print(f"Successfully loaded model from: {try_path}")
                    break
                except Exception as e2:
                    last_error = e2
                    try:
                        self.smpl_model = smplx.create(
                            model_path=try_path,
                            model_type='smpl',
                            gender=gender,
                            batch_size=1
                        ).to(self.device)
                        print(f"Successfully loaded model from: {try_path}")
                        break
                    except Exception as e3:
                        last_error = e3
                        continue
        else:
            error_msg = str(last_error) if last_error else "Unknown error"
            print(f"Error details: {error_msg}")
            raise RuntimeError(
                f"Failed to load SMPL model after trying paths: {model_paths_to_try}. "
                f"Error: {error_msg}. "
                f"Models should be in a 'models' subdirectory. "
                f"Expected files: basicModel_f_lbs_*.pkl (female) or basicmodel_m_lbs_*.pkl (male)"
            )
    
    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 global_orient is None:
            global_orient = torch.zeros([batch_size, 3], device=self.device)
            global_orient[0, 0] = np.radians(2)
        elif isinstance(global_orient, np.ndarray):
            global_orient = torch.FloatTensor(global_orient).to(self.device)
        
        if body_pose is None:
            body_pose = torch.zeros([batch_size, 69], device=self.device)
            
            shoulder_down = np.radians(-12.5)
            shoulder_forward = np.radians(7.5)
            upper_arm_adduction = np.radians(12.5)
            upper_arm_forward = np.radians(7.5)
            elbow_bend = np.radians(12.5)
            palm_inward = np.radians(15)
            hip_forward_tilt = np.radians(2)
            hip_outward = np.radians(7.5)
            hip_flex = np.radians(3.5)
            knee_bend = np.radians(4)
            foot_outward = np.radians(11.5)
            
            body_pose[0, 6:9] = torch.tensor([shoulder_down, 0, shoulder_forward], device=self.device)
            body_pose[0, 9:12] = torch.tensor([shoulder_down, 0, -shoulder_forward], device=self.device)
            
            body_pose[0, 12:15] = torch.tensor([upper_arm_adduction, upper_arm_forward, 0], device=self.device)
            body_pose[0, 15:18] = torch.tensor([-upper_arm_adduction, upper_arm_forward, 0], device=self.device)
            
            body_pose[0, 18:21] = torch.tensor([0, elbow_bend, 0], device=self.device)
            body_pose[0, 21:24] = torch.tensor([0, elbow_bend, 0], device=self.device)
            
            body_pose[0, 24:27] = torch.tensor([0, 0, palm_inward], device=self.device)
            body_pose[0, 27:30] = torch.tensor([0, 0, -palm_inward], device=self.device)
            
            body_pose[0, 30:33] = torch.tensor([np.radians(5), 0, 0], device=self.device)
            body_pose[0, 33:36] = torch.tensor([np.radians(3), 0, 0], device=self.device)
            body_pose[0, 36:39] = torch.tensor([0, 0, 0], device=self.device)
            
            body_pose[0, 39:42] = torch.tensor([np.radians(2), 0, 0], device=self.device)
            body_pose[0, 42:45] = torch.tensor([0, 0, 0], device=self.device)
            
            body_pose[0, 45:48] = torch.tensor([hip_flex, hip_outward, 0], device=self.device)
            body_pose[0, 48:51] = torch.tensor([hip_flex, -hip_outward, 0], device=self.device)
            
            body_pose[0, 51:54] = torch.tensor([0, knee_bend, 0], device=self.device)
            body_pose[0, 54:57] = torch.tensor([0, knee_bend, 0], device=self.device)
            
            body_pose[0, 57:60] = torch.tensor([0, foot_outward, 0], device=self.device)
            body_pose[0, 60:63] = torch.tensor([0, -foot_outward, 0], device=self.device)
            
            body_pose[0, 63:66] = torch.tensor([0, 0, 0], device=self.device)
            body_pose[0, 66:69] = torch.tensor([0, 0, 0], device=self.device)
        elif isinstance(body_pose, np.ndarray):
            body_pose = torch.FloatTensor(body_pose).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 = "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 = "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)