Update src/generate_utils/generate.py

src/generate_utils/generate.py

@@ -26,8 +26,8 @@ class PressureGenerator:
     }
 
     def __init__(self, 
-                 ckpt_dir="src/generate_utils/lib/ckpt/pressurepose_20251222_180032",
-                 smpl_model_dir="src/smpl_models",
+                 ckpt_dir="generate_utils/lib/ckpt/pressurepose_20251222_180032",
+                 smpl_model_dir="E:/pyku/smpl_models",
                  device="cpu"):
         """
         初始化生成器：加载配置、权重和 SMPL 模型。
@@ -58,13 +58,7 @@ class PressureGenerator:
     def _load_config(self):
         config_path = os.path.join(self.ckpt_dir, 'config.yaml')
         if not os.path.exists(config_path):
-            from huggingface_hub import snapshot_download
-            snapshot_download(
-                "yolozyk/PaGe",
-                local_dir="src/generate_utils/lib/ckpt/",
-                local_dir_use_symlinks=False,
-                ignore_patterns=["*.safetensors", ".gitattributes"],
-            )
+            raise FileNotFoundError(f"Config not found at {config_path}")
         with open(config_path, 'r') as f:
             return yaml.safe_load(f)
 
@@ -73,13 +67,7 @@ class PressureGenerator:
         ckpt_path = os.path.join(self.ckpt_dir, 'ckpts', 'best_model.pth')
         
         if not os.path.exists(ckpt_path):
-            from huggingface_hub import snapshot_download
-            snapshot_download(
-                "yolozyk/PaGe",
-                local_dir="src/generate_utils/lib/ckpt/",
-                local_dir_use_symlinks=False,
-                ignore_patterns=["*.safetensors", ".gitattributes"],
-            )
+            raise FileNotFoundError(f"Checkpoint not found at {ckpt_path}")
 
         checkpoint = torch.load(ckpt_path, map_location=self.device)
         model.load_state_dict(checkpoint['model_state_dict'])
@@ -97,7 +85,7 @@ class PressureGenerator:
         return smpl
 
     @torch.no_grad()
-    def generate(self, betas, transl, poses):
+    def generate(self, betas, transl, poses, transfer=False):
         """
         执行推理。
         输入参数应该是 Tensor, 维度需符合模型要求 (Batch Size, ...)。
@@ -116,6 +104,10 @@ class PressureGenerator:
         )
         vertices = output.vertices
 
+        if transfer:
+            vertices[:, :, 1] = 1.80 - vertices[:, :, 1]
+            vertices[:, :, 2] = -vertices[:, :, 2]
+
         # 2. 预测压力图
         pred_pmap = self.cvae_model.inference(vertices)