add smpl

smpl_generator.py

@@ -23,15 +23,10 @@ class SMPLGenerator:
             else:
                 model_path_obj.mkdir(parents=True, exist_ok=True)
         
-        models_source = model_path_obj / "smpl" / "models"
+        models_source = Path("smpl/smpl/models")
         if not models_source.exists():
             models_source = model_path_obj / "models"
         
-        if models_source.exists() and (model_path_obj / "smpl").exists():
-            expected_models_dir = model_path_obj / "smpl" / "models"
-            if not expected_models_dir.exists() and models_source.exists():
-                print(f"Models found in {models_source}, smplx may need them in {expected_models_dir}")
-        
         self.model_path = model_path_obj
         model_path_str = str(self.model_path)
         
@@ -39,21 +34,32 @@ class SMPLGenerator:
             gender = "male"
             print("Note: Neutral gender not available, using male model")
         
-        models_dir = model_path_obj / "models"
+        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]}")
+            
+            expected_smpl_dir = model_path_obj / "smpl"
+            expected_models_dir = expected_smpl_dir / "models"
+            
+            if not expected_models_dir.exists() and models_source.exists():
+                print(f"Creating expected directory structure: {expected_models_dir}")
+                expected_models_dir.mkdir(parents=True, exist_ok=True)
+                
+                for model_file in model_files:
+                    target_file = expected_models_dir / model_file.name
+                    if not target_file.exists():
+                        import shutil
+                        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 / "smpl" / "models"
-        
-        print(f"Looking for SMPL models in: {model_path_str}")
-        if models_dir.exists():
-            model_files = list(models_dir.glob("*.pkl"))
-            print(f"Found {len(model_files)} model files: {[f.name for f in model_files]}")
+            models_dir = model_path_obj / "models"
         
         model_paths_to_try = [
+            str(model_path_obj),
             "smpl",
-            model_path_str,
-            str(model_path_obj / "smpl"),
-            str(models_dir.parent) if models_dir.exists() else None,
-            "smpl/smpl"
+            str(models_dir.parent) if models_dir.exists() else None
         ]
         model_paths_to_try = [p for p in model_paths_to_try if p is not None and Path(p).exists()]