| 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) | |