update generate tools

src/generate_utils/__init__.py

@@ -0,0 +1 @@
+from .generate import PressureGenerator

src/generate_utils/generate.py

@@ -0,0 +1,118 @@
+import os
+import yaml
+import smplx
+import torch
+
+from .lib.model.cvae import SMPL2PressureCVAE
+
+class PressureGenerator:
+    # 静态配置数据
+    DATASET_META = {
+        'tip': {
+            'max_p': 512.0,
+            'crop_size': [56, 40], 
+            'path': "/workspace/zyk/public_data/wzy_opt_dataset_w_feats"
+        }, 
+        'pressurepose': {
+            'max_p': 100.0, 
+            'crop_size': [64, 27], 
+            'path': "/workspace/zyk/public_data/pressurepose/synth"
+        }, 
+        'moyo': {
+            'max_p': 64.0, 
+            'crop_size': [110, 37], 
+            'path': "/workspace/zyk/public_data/moyo"
+        }
+    }
+
+    def __init__(self, 
+                 ckpt_dir="generate_utils/lib/ckpt/pressurepose_20251222_180032",
+                 smpl_model_dir="E:/pyku/smpl_models",
+                 device="cpu"):
+        """
+        初始化生成器：加载配置、权重和 SMPL 模型。
+        """
+        self.device = torch.device(device)
+        self.ckpt_dir = ckpt_dir
+        self.smpl_model_dir = smpl_model_dir
+
+        # 1. 加载配置
+        self.cfg = self._load_config()
+        
+        # 2. 设置数据集相关参数
+        dataset_name = self.cfg['dataset']['name']
+        if dataset_name not in self.DATASET_META:
+             raise ValueError(f"Unknown dataset name: {dataset_name}")
+             
+        self.max_pressure = self.DATASET_META[dataset_name]['max_p']
+        self.is_normalized = self.cfg['dataset'].get('normal', False)
+
+        # 3. 加载 CVAE 模型
+        self.cvae_model = self._load_cvae()
+
+        # 4. 加载 SMPL 模型
+        self.smpl_model = self._load_smpl()
+        
+        print(f"Pressure Generator loaded successfully from {self.ckpt_dir}")
+
+    def _load_config(self):
+        config_path = os.path.join(self.ckpt_dir, 'config.yaml')
+        if not os.path.exists(config_path):
+            raise FileNotFoundError(f"Config not found at {config_path}")
+        with open(config_path, 'r') as f:
+            return yaml.safe_load(f)
+
+    def _load_cvae(self):
+        model = SMPL2PressureCVAE(self.cfg).to(self.device)
+        ckpt_path = os.path.join(self.ckpt_dir, 'ckpts', 'best_model.pth')
+        
+        if not os.path.exists(ckpt_path):
+            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'])
+        model.eval()
+        return model
+
+    def _load_smpl(self):
+        # 创建 SMPL 模型 (这是一个比较耗时的操作)
+        smpl = smplx.create(
+            self.smpl_model_dir, 
+            model_type='smpl', 
+            gender='neutral', 
+            ext='pkl'
+        ).to(self.device)
+        return smpl
+
+    @torch.no_grad()
+    def generate(self, betas, transl, poses):
+        """
+        执行推理。
+        输入参数应该是 Tensor, 维度需符合模型要求 (Batch Size, ...)。
+        """
+        # 确保输入在正确的设备上
+        if betas.device != self.device: betas = betas.to(self.device)
+        if transl.device != self.device: transl = transl.to(self.device)
+        if poses.device != self.device: poses = poses.to(self.device)
+
+        # 1. 获取 SMPL 顶点 (Vertices)
+        output = self.smpl_model(
+            betas=betas,
+            global_orient=poses[:, :3], # 前3位是全局旋转
+            body_pose=poses[:, 3:],     # 后69位是身体姿态
+            transl=transl,
+        )
+        vertices = output.vertices
+
+        # 2. 预测压力图
+        pred_pmap = self.cvae_model.inference(vertices)
+
+        # 3. 后处理 (反归一化 & 阈值过滤)
+        if self.is_normalized:
+            pred_pmap = pred_pmap * self.max_pressure
+        
+        # 这里的 0.1 是硬编码的阈值，也可以提取为参数
+        pred_pmap[pred_pmap < 0.1] = 0
+
+        return pred_pmap
+

src/generate_utils/lib/ckpt/pressurepose_20251222_180032/config.yaml

@@ -0,0 +1,40 @@
+dataset:
+  name: pressurepose
+  normal: true
+device: cuda
+model:
+  cond_decoder:
+    type: mlp
+  cond_embed_dim: 256
+  cond_encoder:
+    return_global_feature: true
+    type: pointnet
+    use_spatial_transformer: false
+  cond_type: vertices
+  dropout_rate: 0.3
+  embed_dim: 256
+  main_decoder:
+    bilinear: false
+    type: unet
+  main_encoder:
+    type: resnet18
+output: null
+seed: 42
+training:
+  batch_size: 64
+  learning_rate: 0.001
+  log_freq: 10
+  loss:
+    kl_weight: 10.0
+    pcr_weight: 60.0
+    pmr_weight: 100.0
+  num_epochs: 100
+  optimizer:
+    type: adam
+    weight_decay: 0.0001
+  save_freq: 999
+  scheduler:
+    gamma: 0.1
+    step_size: 30
+    type: cosine
+  val_freq: 1

src/generate_utils/lib/ckpt/pressurepose_20251222_180032/logs/train.log

@@ -0,0 +1,153 @@
+2025-12-22 18:00:32,951 - INFO - Starting training on cuda:0: pressurepose...
+2025-12-22 18:08:37,856 - INFO - Epoch 0 | Train Loss: 3.9866 | Val Loss: 2.2416
+2025-12-22 18:08:47,179 - INFO - --> Best model saved at epoch 0
+2025-12-22 18:16:45,378 - INFO - Epoch 1 | Train Loss: 2.1470 | Val Loss: 2.0693
+2025-12-22 18:16:54,269 - INFO - --> Best model saved at epoch 1
+2025-12-22 18:24:54,079 - INFO - Epoch 2 | Train Loss: 1.9345 | Val Loss: 1.7783
+2025-12-22 18:25:02,726 - INFO - --> Best model saved at epoch 2
+2025-12-22 18:32:55,954 - INFO - Epoch 3 | Train Loss: 1.7918 | Val Loss: 1.5810
+2025-12-22 18:33:03,955 - INFO - --> Best model saved at epoch 3
+2025-12-22 18:41:00,038 - INFO - Epoch 4 | Train Loss: 1.7060 | Val Loss: 1.5400
+2025-12-22 18:41:06,723 - INFO - --> Best model saved at epoch 4
+2025-12-22 18:49:04,843 - INFO - Epoch 5 | Train Loss: 1.6220 | Val Loss: 1.5549
+2025-12-22 18:57:05,040 - INFO - Epoch 6 | Train Loss: 1.5575 | Val Loss: 1.3848
+2025-12-22 18:57:11,939 - INFO - --> Best model saved at epoch 6
+2025-12-22 19:05:07,141 - INFO - Epoch 7 | Train Loss: 1.5061 | Val Loss: 1.3306
+2025-12-22 19:05:17,022 - INFO - --> Best model saved at epoch 7
+2025-12-22 19:13:12,365 - INFO - Epoch 8 | Train Loss: 1.4657 | Val Loss: 1.3127
+2025-12-22 19:13:21,550 - INFO - --> Best model saved at epoch 8
+2025-12-22 19:21:14,461 - INFO - Epoch 9 | Train Loss: 1.4281 | Val Loss: 1.2233
+2025-12-22 19:21:22,432 - INFO - --> Best model saved at epoch 9
+2025-12-22 19:29:17,334 - INFO - Epoch 10 | Train Loss: 1.3971 | Val Loss: 1.3038
+2025-12-22 19:37:16,753 - INFO - Epoch 11 | Train Loss: 1.3737 | Val Loss: 1.2679
+2025-12-22 19:45:15,562 - INFO - Epoch 12 | Train Loss: 1.3458 | Val Loss: 1.1456
+2025-12-22 19:45:22,194 - INFO - --> Best model saved at epoch 12
+2025-12-22 19:53:23,432 - INFO - Epoch 13 | Train Loss: 1.3215 | Val Loss: 1.1579
+2025-12-22 20:01:26,931 - INFO - Epoch 14 | Train Loss: 1.2890 | Val Loss: 1.1414
+2025-12-22 20:01:35,934 - INFO - --> Best model saved at epoch 14
+2025-12-22 20:09:33,396 - INFO - Epoch 15 | Train Loss: 1.2620 | Val Loss: 1.0436
+2025-12-22 20:09:42,377 - INFO - --> Best model saved at epoch 15
+2025-12-22 20:17:41,945 - INFO - Epoch 16 | Train Loss: 1.2347 | Val Loss: 1.0598
+2025-12-22 20:26:25,986 - INFO - Epoch 17 | Train Loss: 1.2038 | Val Loss: 1.1026
+2025-12-22 20:35:12,185 - INFO - Epoch 18 | Train Loss: 1.1821 | Val Loss: 0.9907
+2025-12-22 20:35:19,867 - INFO - --> Best model saved at epoch 18
+2025-12-22 20:44:01,826 - INFO - Epoch 19 | Train Loss: 1.1594 | Val Loss: 0.9307
+2025-12-22 20:44:11,053 - INFO - --> Best model saved at epoch 19
+2025-12-22 20:52:25,588 - INFO - Epoch 20 | Train Loss: 1.1332 | Val Loss: 0.8907
+2025-12-22 20:52:33,744 - INFO - --> Best model saved at epoch 20
+2025-12-22 21:00:36,050 - INFO - Epoch 21 | Train Loss: 1.1235 | Val Loss: 0.8508
+2025-12-22 21:00:42,396 - INFO - --> Best model saved at epoch 21
+2025-12-22 21:08:38,121 - INFO - Epoch 22 | Train Loss: 1.1021 | Val Loss: 0.9147
+2025-12-22 21:16:36,098 - INFO - Epoch 23 | Train Loss: 1.0872 | Val Loss: 0.9329
+2025-12-22 21:24:32,846 - INFO - Epoch 24 | Train Loss: 1.0739 | Val Loss: 0.8428
+2025-12-22 21:24:40,586 - INFO - --> Best model saved at epoch 24
+2025-12-22 21:32:41,914 - INFO - Epoch 25 | Train Loss: 1.0588 | Val Loss: 0.8102
+2025-12-22 21:32:51,039 - INFO - --> Best model saved at epoch 25
+2025-12-22 21:40:50,536 - INFO - Epoch 26 | Train Loss: 1.0506 | Val Loss: 0.8025
+2025-12-22 21:40:59,505 - INFO - --> Best model saved at epoch 26
+2025-12-22 21:48:55,531 - INFO - Epoch 27 | Train Loss: 1.0417 | Val Loss: 0.7854
+2025-12-22 21:49:03,636 - INFO - --> Best model saved at epoch 27
+2025-12-22 21:57:02,459 - INFO - Epoch 28 | Train Loss: 1.0373 | Val Loss: 0.7912
+2025-12-22 22:05:07,235 - INFO - Epoch 29 | Train Loss: 1.0243 | Val Loss: 0.7809
+2025-12-22 22:05:17,305 - INFO - --> Best model saved at epoch 29
+2025-12-22 22:13:14,376 - INFO - Epoch 30 | Train Loss: 1.0116 | Val Loss: 0.8382
+2025-12-22 22:21:12,713 - INFO - Epoch 31 | Train Loss: 1.0017 | Val Loss: 0.7629
+2025-12-22 22:21:21,674 - INFO - --> Best model saved at epoch 31
+2025-12-22 22:29:18,931 - INFO - Epoch 32 | Train Loss: 0.9948 | Val Loss: 0.8141
+2025-12-22 22:37:20,810 - INFO - Epoch 33 | Train Loss: 0.9933 | Val Loss: 0.7737
+2025-12-22 22:45:23,264 - INFO - Epoch 34 | Train Loss: 0.9768 | Val Loss: 0.7786
+2025-12-22 22:53:25,925 - INFO - Epoch 35 | Train Loss: 0.9706 | Val Loss: 0.7660
+2025-12-22 23:01:28,145 - INFO - Epoch 36 | Train Loss: 0.9622 | Val Loss: 0.8475
+2025-12-22 23:09:31,131 - INFO - Epoch 37 | Train Loss: 0.9572 | Val Loss: 0.7795
+2025-12-22 23:17:32,236 - INFO - Epoch 38 | Train Loss: 0.9484 | Val Loss: 0.7872
+2025-12-22 23:25:34,544 - INFO - Epoch 39 | Train Loss: 0.9437 | Val Loss: 0.7449
+2025-12-22 23:25:43,984 - INFO - --> Best model saved at epoch 39
+2025-12-22 23:33:42,019 - INFO - Epoch 40 | Train Loss: 0.9383 | Val Loss: 0.7207
+2025-12-22 23:33:51,375 - INFO - --> Best model saved at epoch 40
+2025-12-22 23:41:48,820 - INFO - Epoch 41 | Train Loss: 0.9304 | Val Loss: 0.7335
+2025-12-22 23:49:51,966 - INFO - Epoch 42 | Train Loss: 0.9212 | Val Loss: 0.8034
+2025-12-22 23:57:56,083 - INFO - Epoch 43 | Train Loss: 0.9181 | Val Loss: 0.7071
+2025-12-22 23:58:05,215 - INFO - --> Best model saved at epoch 43
+2025-12-23 00:06:03,136 - INFO - Epoch 44 | Train Loss: 0.9082 | Val Loss: 0.7624
+2025-12-23 00:14:12,478 - INFO - Epoch 45 | Train Loss: 0.9088 | Val Loss: 0.7162
+2025-12-23 00:22:18,581 - INFO - Epoch 46 | Train Loss: 0.8930 | Val Loss: 0.7205
+2025-12-23 00:30:23,463 - INFO - Epoch 47 | Train Loss: 0.8928 | Val Loss: 0.7508
+2025-12-23 00:38:20,809 - INFO - Epoch 48 | Train Loss: 0.8809 | Val Loss: 0.6766
+2025-12-23 00:38:27,973 - INFO - --> Best model saved at epoch 48
+2025-12-23 00:46:29,782 - INFO - Epoch 49 | Train Loss: 0.8810 | Val Loss: 0.6758
+2025-12-23 00:46:38,221 - INFO - --> Best model saved at epoch 49
+2025-12-23 00:54:38,092 - INFO - Epoch 50 | Train Loss: 0.8696 | Val Loss: 0.7521
+2025-12-23 01:03:00,051 - INFO - Epoch 51 | Train Loss: 0.8633 | Val Loss: 0.6730
+2025-12-23 01:03:07,488 - INFO - --> Best model saved at epoch 51
+2025-12-23 01:11:19,673 - INFO - Epoch 52 | Train Loss: 0.8607 | Val Loss: 0.6512
+2025-12-23 01:11:27,660 - INFO - --> Best model saved at epoch 52
+2025-12-23 01:19:24,289 - INFO - Epoch 53 | Train Loss: 0.8507 | Val Loss: 0.6400
+2025-12-23 01:19:33,796 - INFO - --> Best model saved at epoch 53
+2025-12-23 01:27:29,840 - INFO - Epoch 54 | Train Loss: 0.8456 | Val Loss: 0.6368
+2025-12-23 01:27:37,645 - INFO - --> Best model saved at epoch 54
+2025-12-23 01:35:32,997 - INFO - Epoch 55 | Train Loss: 0.8410 | Val Loss: 0.6299
+2025-12-23 01:35:42,289 - INFO - --> Best model saved at epoch 55
+2025-12-23 01:43:44,539 - INFO - Epoch 56 | Train Loss: 0.8303 | Val Loss: 0.6753
+2025-12-23 01:52:07,397 - INFO - Epoch 57 | Train Loss: 0.8289 | Val Loss: 0.6344
+2025-12-23 02:00:30,969 - INFO - Epoch 58 | Train Loss: 0.8215 | Val Loss: 0.6222
+2025-12-23 02:00:39,050 - INFO - --> Best model saved at epoch 58
+2025-12-23 02:08:56,943 - INFO - Epoch 59 | Train Loss: 0.8130 | Val Loss: 0.6418
+2025-12-23 02:17:16,645 - INFO - Epoch 60 | Train Loss: 0.8074 | Val Loss: 0.6353
+2025-12-23 02:25:12,935 - INFO - Epoch 61 | Train Loss: 0.8003 | Val Loss: 0.6470
+2025-12-23 02:33:12,804 - INFO - Epoch 62 | Train Loss: 0.7943 | Val Loss: 0.6385
+2025-12-23 02:41:13,969 - INFO - Epoch 63 | Train Loss: 0.7891 | Val Loss: 0.5937
+2025-12-23 02:41:22,714 - INFO - --> Best model saved at epoch 63
+2025-12-23 02:49:17,641 - INFO - Epoch 64 | Train Loss: 0.7824 | Val Loss: 0.5996
+2025-12-23 02:57:15,236 - INFO - Epoch 65 | Train Loss: 0.7795 | Val Loss: 0.5926
+2025-12-23 02:57:22,741 - INFO - --> Best model saved at epoch 65
+2025-12-23 03:05:19,890 - INFO - Epoch 66 | Train Loss: 0.7695 | Val Loss: 0.5950
+2025-12-23 03:13:22,173 - INFO - Epoch 67 | Train Loss: 0.7681 | Val Loss: 0.6047
+2025-12-23 03:21:19,313 - INFO - Epoch 68 | Train Loss: 0.7563 | Val Loss: 0.5717
+2025-12-23 03:21:26,615 - INFO - --> Best model saved at epoch 68
+2025-12-23 03:29:19,295 - INFO - Epoch 69 | Train Loss: 0.7572 | Val Loss: 0.5887
+2025-12-23 03:37:18,966 - INFO - Epoch 70 | Train Loss: 0.7479 | Val Loss: 0.6256
+2025-12-23 03:45:16,737 - INFO - Epoch 71 | Train Loss: 0.7422 | Val Loss: 0.5700
+2025-12-23 03:45:25,848 - INFO - --> Best model saved at epoch 71
+2025-12-23 03:53:20,776 - INFO - Epoch 72 | Train Loss: 0.7358 | Val Loss: 0.5686
+2025-12-23 03:53:30,797 - INFO - --> Best model saved at epoch 72
+2025-12-23 04:01:26,968 - INFO - Epoch 73 | Train Loss: 0.7304 | Val Loss: 0.5606
+2025-12-23 04:01:33,584 - INFO - --> Best model saved at epoch 73
+2025-12-23 04:09:25,824 - INFO - Epoch 74 | Train Loss: 0.7257 | Val Loss: 0.5824
+2025-12-23 04:17:23,933 - INFO - Epoch 75 | Train Loss: 0.7163 | Val Loss: 0.5535
+2025-12-23 04:17:32,977 - INFO - --> Best model saved at epoch 75
+2025-12-23 04:25:27,001 - INFO - Epoch 76 | Train Loss: 0.7151 | Val Loss: 0.5496
+2025-12-23 04:25:34,823 - INFO - --> Best model saved at epoch 76
+2025-12-23 04:33:28,463 - INFO - Epoch 77 | Train Loss: 0.7059 | Val Loss: 0.5530
+2025-12-23 04:41:27,898 - INFO - Epoch 78 | Train Loss: 0.7017 | Val Loss: 0.5554
+2025-12-23 04:49:26,990 - INFO - Epoch 79 | Train Loss: 0.6985 | Val Loss: 0.5580
+2025-12-23 04:57:24,784 - INFO - Epoch 80 | Train Loss: 0.6916 | Val Loss: 0.5468
+2025-12-23 04:57:34,075 - INFO - --> Best model saved at epoch 80
+2025-12-23 05:05:27,512 - INFO - Epoch 81 | Train Loss: 0.6877 | Val Loss: 0.5476
+2025-12-23 05:13:25,319 - INFO - Epoch 82 | Train Loss: 0.6840 | Val Loss: 0.5423
+2025-12-23 05:13:32,737 - INFO - --> Best model saved at epoch 82
+2025-12-23 05:21:35,339 - INFO - Epoch 83 | Train Loss: 0.6788 | Val Loss: 0.5242
+2025-12-23 05:21:43,953 - INFO - --> Best model saved at epoch 83
+2025-12-23 05:30:01,032 - INFO - Epoch 84 | Train Loss: 0.6723 | Val Loss: 0.5558
+2025-12-23 05:38:07,292 - INFO - Epoch 85 | Train Loss: 0.6695 | Val Loss: 0.5458
+2025-12-23 05:46:03,208 - INFO - Epoch 86 | Train Loss: 0.6647 | Val Loss: 0.5323
+2025-12-23 05:54:00,838 - INFO - Epoch 87 | Train Loss: 0.6594 | Val Loss: 0.5247
+2025-12-23 06:01:57,323 - INFO - Epoch 88 | Train Loss: 0.6564 | Val Loss: 0.5216
+2025-12-23 06:02:06,118 - INFO - --> Best model saved at epoch 88
+2025-12-23 06:09:58,071 - INFO - Epoch 89 | Train Loss: 0.6512 | Val Loss: 0.5181
+2025-12-23 06:10:07,271 - INFO - --> Best model saved at epoch 89
+2025-12-23 06:18:02,564 - INFO - Epoch 90 | Train Loss: 0.6501 | Val Loss: 0.5153
+2025-12-23 06:18:11,750 - INFO - --> Best model saved at epoch 90
+2025-12-23 06:26:04,938 - INFO - Epoch 91 | Train Loss: 0.6478 | Val Loss: 0.5166
+2025-12-23 06:34:05,888 - INFO - Epoch 92 | Train Loss: 0.6415 | Val Loss: 0.5184
+2025-12-23 06:42:06,933 - INFO - Epoch 93 | Train Loss: 0.6407 | Val Loss: 0.5158
+2025-12-23 06:50:05,590 - INFO - Epoch 94 | Train Loss: 0.6373 | Val Loss: 0.5138
+2025-12-23 06:50:14,737 - INFO - --> Best model saved at epoch 94
+2025-12-23 06:58:09,245 - INFO - Epoch 95 | Train Loss: 0.6369 | Val Loss: 0.5102
+2025-12-23 06:58:16,989 - INFO - --> Best model saved at epoch 95
+2025-12-23 07:06:09,566 - INFO - Epoch 96 | Train Loss: 0.6362 | Val Loss: 0.5114
+2025-12-23 07:14:08,758 - INFO - Epoch 97 | Train Loss: 0.6331 | Val Loss: 0.5140
+2025-12-23 07:22:02,795 - INFO - Epoch 98 | Train Loss: 0.6318 | Val Loss: 0.5099
+2025-12-23 07:22:11,592 - INFO - --> Best model saved at epoch 98
+2025-12-23 07:30:05,162 - INFO - Epoch 99 | Train Loss: 0.6318 | Val Loss: 0.5092
+2025-12-23 07:30:12,362 - INFO - --> Best model saved at epoch 99
+2025-12-23 07:30:12,362 - INFO - Training complete.

src/generate_utils/lib/ckpt/pressurepose_20251222_180032/test_result.txt

@@ -0,0 +1,20 @@
+Test Results for experiment: output/pressurepose_20251222_180032/
+Checkpoint used: best_model.pth
+------------------------------
+SSIM: 0.892187
+MAE: 1.883694
+MSE: 31.912682
+MRE: 0.018837
+CoP_Dist: 0.500346
+------------------------------
+IPMAN SSIM: 0.527497
+IPMAN MAE: 5.092346
+IPMAN MSE: 221.172048
+IPMAN MRE: 0.050923
+IPMAN CoP_Dist: 1.602944
+------------------------------
+PMR SSIM: 0.541821
+PMR MAE: 16.486778
+PMR MSE: 2864.192610
+PMR MRE: 0.164868
+PMR CoP_Dist: 4.019547

src/generate_utils/lib/model/cvae.py

@@ -0,0 +1,180 @@
+import os
+import torch
+import torch.nn as nn
+from .resnet import resnet18, resnet34, resnet50
+from .unet import UNetEncoder, UNetDecoder
+from .pointnet import create_pointnet_encoder
+from .mlp import create_pointcloud_decoder
+
+DATASET_META = {
+    'tip': {
+        'max_p': 512.0,
+        'crop_size': [56, 40], 
+        'path': "/workspace/zyk/public_data/wzy_opt_dataset_w_feats"
+    }, 
+    'pressurepose': {
+        'max_p': 100.0, 
+        'crop_size': [64, 27], 
+        'path': "/workspace/zyk/public_data/pressurepose/synth"
+    }, 
+    'moyo': {
+        'max_p': 64.0, 
+        'crop_size': [110, 37], 
+        'path': "/workspace/zyk/public_data/moyo"
+    }
+}
+
+
+
+class SMPL2PressureCVAE(nn.Module):
+    """
+    SMPL2Pressure cVAE model with dual branches.
+    Main Branch: Pressure map encoding and reconstruction.
+    Condition Branch: Point cloud encoding and reconstruction.
+    """
+    def __init__(self, cfg):
+        super(SMPL2PressureCVAE, self).__init__()
+        self.cfg = cfg
+        self.embed_dim = cfg['model']['embed_dim']
+        self.cond_embed_dim = cfg['model']['cond_embed_dim']
+        
+        # 1. Main Encoder (Pressure Map -> z_params)
+        # Supports ResNet or UNet as the visual encoder
+        main_enc_type = cfg['model']['main_encoder']['type']
+        if "resnet" in main_enc_type:
+            # Using the ResNet implementation provided
+            model_func = eval(main_enc_type)
+            self.main_encoder = model_func(
+                embed_dim=self.embed_dim, 
+                cond_embed_dim=self.cond_embed_dim,
+                dp_rate=cfg['model']['dropout_rate']
+            )
+        else:
+            self.main_encoder = UNetEncoder(
+                cond_dim=self.cond_embed_dim,
+                embed_dim=self.embed_dim,
+                crop=DATASET_META[cfg['dataset']['name']]['crop_size'] # Need to pass this from config
+            )
+
+        # 2. Condition Encoder (Point Cloud -> cond_features)
+        self.cond_encoder = create_pointnet_encoder(
+            input_dim=3,
+            feature_dim=self.cond_embed_dim,
+            use_spatial_transformer=cfg['model']['cond_encoder']['use_spatial_transformer'],
+            return_global_feature=cfg['model']['cond_encoder']['return_global_feature']
+        )
+
+        # 3. Main Decoder (z + cond -> Pressure Map)
+        self.main_decoder = UNetDecoder(
+            cond_dim=self.cond_embed_dim,
+            embed_dim=self.embed_dim,
+            bilinear=cfg['model']['main_decoder']['bilinear'],
+            crop=DATASET_META[cfg['dataset']['name']]['crop_size']
+        )
+
+        # 4. Condition Decoder (cond_features -> Point Cloud)
+        self.cond_decoder = create_pointcloud_decoder(
+            latent_dim=self.embed_dim,
+            num_points=6890, # SMPL vertices
+            architecture=cfg['model']['cond_decoder']['type']
+        )
+
+    def load_pretrained_cond(self, path):
+        """
+        专门用于加载预训练的点云编解码分支权重
+        """
+        if not os.path.exists(path):
+            print(f"Warning: Pretrained path {path} not found. Training from scratch.")
+            return
+
+        print(f"Loading pretrained condition branch from {path}...")
+        ckpt = torch.load(path, map_location='cpu')
+        
+        # 加载 encoder 和 decoder
+        self.cond_encoder.load_state_dict(ckpt['cond_encoder'])
+        self.cond_decoder.load_state_dict(ckpt['cond_decoder'])
+        
+        # 可选：如果你希望预训练的组件在主训练初期不被破坏，可以冻结它们
+        # for param in self.cond_encoder.parameters(): param.requires_grad = False
+
+    def reparameterize(self, mu, log_var):
+        """Reparameterization trick to sample z from N(mu, var)"""
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return mu + eps * std
+
+    def forward(self, pressure_map, vertices):
+        """
+        Forward pass during training.
+        Args:
+            pressure_map: (B, 1, H, W)
+            vertices: (B, 6890, 3)
+        """
+        # A. Encode Point Cloud to get condition (Condition branch)
+        # PointNet expects (B, 3, N)
+        pts = vertices.transpose(2, 1)
+        cond_feat = self.cond_encoder(pts) # (B, cond_embed_dim)
+
+        # B. Encode Pressure Map with condition to get latent distribution
+        # Note: resnet implementation already handles concatenation inside
+        mu, log_var = self.main_encoder(pressure_map, cond_feat)
+        
+        # C. Sample z
+        z = self.reparameterize(mu, log_var)
+
+        # D. Reconstruct Pressure Map
+        recon_pressure = self.main_decoder(z, cond_feat)
+
+        # E. Reconstruct Point Cloud (Side task for better latent space)
+        recon_vertices = self.cond_decoder(cond_feat)
+
+        return {
+            'recon_pressure': recon_pressure,
+            'recon_vertices': recon_vertices,
+            'mu': mu,
+            'log_var': log_var,
+            'z': z
+        }
+
+    @torch.no_grad()
+    def inference(self, vertices):
+        """
+        Inference: Generate pressure from Point Cloud only.
+        """
+        self.eval()
+        # 1. Encode Point Cloud
+        pts = vertices.transpose(2, 1)
+        cond_feat = self.cond_encoder(pts)
+
+        # 2. Sample z from prior N(0, 1)
+        z = torch.randn(vertices.size(0), self.embed_dim).to(vertices.device)
+
+        # 3. Decode to Pressure Map
+        gen_pressure = self.main_decoder(z, cond_feat)
+        
+        return gen_pressure
+    
+
+if __name__ == "__main__":
+    import yaml
+
+    cfg_path = 'config/config_base.yaml'
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    with open(cfg_path, 'r') as f:
+        cfg = yaml.safe_load(f)
+
+    model = SMPL2PressureCVAE(cfg).to(device)
+
+    pressure = torch.randn(8, 1, 56, 40).to(device)
+    vertices = torch.randn(8, 6890, 3).to(device)
+
+    res = model(pressure, vertices)
+
+    pred_pressure = model.inference(vertices)
+
+    import pdb; pdb.set_trace()
+
+
+
+

src/generate_utils/lib/model/loss.py

@@ -0,0 +1,61 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class SMPL2PressureLoss(nn.Module):
+    """
+    Combined loss for SMPL2Pressure cVAE.
+    - Pressure Map Reconstruction (PMR): MSE Loss
+    - Point Cloud Reconstruction (PCR): MSE Loss (from condition features)
+    - KL Divergence: Regularization of latent space
+    """
+    def __init__(self, cfg):
+        super(SMPL2PressureLoss, self).__init__()
+        self.cfg_loss = cfg['training']['loss']
+        
+        # 权重配置
+        self.pmr_weight = self.cfg_loss.get('pmr_weight', 10.0)
+        self.pcr_weight = self.cfg_loss.get('pcr_weight', 6.0)
+        self.kl_weight = self.cfg_loss.get('kl_weight', 2.0)
+
+    def forward(self, outputs, target_pressure, target_vertices):
+        """
+        Args:
+            outputs: cVAE模型的输出字典
+            target_pressure: 真值压力图 (B, H, W) 或 (B, 1, H, W)
+            target_vertices: 真值SMPL顶点 (B, 6890, 3)
+        """
+        # 1. 压力图重建损失 (PMR)
+        recon_pressure = outputs['recon_pressure']
+        # 统一维度: 确保 target 也是 (B, 1, H, W)
+        if recon_pressure.dim() == 3:
+            recon_pressure = recon_pressure.unsqueeze(1)
+
+        if target_pressure.dim() == 3:
+            target_pressure = target_pressure.unsqueeze(1)
+        
+        loss_pmr = F.mse_loss(recon_pressure, target_pressure, reduction='mean')
+
+        # 2. 点云重建损失 (PCR)
+        # 根据你的提醒，这里的 recon_vertices 应该是从 cond_features 解码出来的
+        recon_vertices = outputs['recon_vertices']
+        loss_pcr = F.mse_loss(recon_vertices, target_vertices, reduction='mean')
+
+        # 3. KL 散度损失
+        mu = outputs['mu']
+        log_var = outputs['log_var']
+        # KL = -0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
+        loss_kl = -0.5 * torch.mean(torch.sum(1 + log_var - mu.pow(2) - log_var.exp(), dim=1))
+
+        # 4. 总损失加权
+        total_loss = (self.pmr_weight * loss_pmr + 
+                      self.pcr_weight * loss_pcr + 
+                      self.kl_weight * loss_kl)
+
+        return {
+            'loss': total_loss,
+            'loss_pmr': loss_pmr,
+            'loss_pcr': loss_pcr,
+            'loss_kl': loss_kl
+        }
+    

src/generate_utils/lib/model/mlp.py

@@ -0,0 +1,382 @@
+"""
+MLP-based decoder for reconstructing point clouds from latent codes.
+
+This module provides flexible MLP architectures for decoding latent representations
+into point clouds, commonly used in generative models like VAE and cVAE.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import List, Optional
+
+
+class PointCloudDecoder(nn.Module):
+    """
+    MLP-based decoder for reconstructing point clouds from latent representations.
+
+    This decoder uses a series of fully connected layers with optional dropout
+    and normalization to transform a latent code into point cloud coordinates.
+
+    Args:
+        latent_dim: Dimensionality of input latent code
+        num_points: Number of output points in the point cloud
+        point_dim: Dimensionality of each point (default: 3 for XYZ coordinates)
+        hidden_dims: List of hidden layer dimensions (default: [1024, 2048])
+        dropout_rate: Dropout probability (default: 0.3)
+        use_batch_norm: Whether to use batch normalization (default: False)
+        activation: Activation function to use (default: 'relu')
+
+    Input:
+        Latent code of shape (B, latent_dim)
+
+    Output:
+        Point cloud of shape (B, num_points, point_dim)
+    """
+
+    def __init__(
+        self,
+        latent_dim: int,
+        num_points: int,
+        point_dim: int = 3,
+        hidden_dims: Optional[List[int]] = None,
+        dropout_rate: float = 0.3,
+        use_batch_norm: bool = False,
+        activation: str = 'relu'
+    ):
+        super(PointCloudDecoder, self).__init__()
+        self.latent_dim = latent_dim
+        self.num_points = num_points
+        self.point_dim = point_dim
+        self.dropout_rate = dropout_rate
+        self.use_batch_norm = use_batch_norm
+
+        # Default hidden dimensions
+        if hidden_dims is None:
+            hidden_dims = [1024, 2048]
+        self.hidden_dims = hidden_dims
+
+        # Select activation function
+        if activation == 'relu':
+            self.activation = nn.ReLU()
+        elif activation == 'leaky_relu':
+            self.activation = nn.LeakyReLU(0.2)
+        elif activation == 'elu':
+            self.activation = nn.ELU()
+        elif activation == 'gelu':
+            self.activation = nn.GELU()
+        else:
+            raise ValueError(f"Unsupported activation: {activation}")
+
+        # Build network layers
+        self.layers = nn.ModuleList()
+        self.batch_norms = nn.ModuleList() if use_batch_norm else None
+        self.dropouts = nn.ModuleList()
+
+        # Input layer
+        prev_dim = latent_dim
+        for hidden_dim in hidden_dims:
+            self.layers.append(nn.Linear(prev_dim, hidden_dim))
+            if use_batch_norm:
+                self.batch_norms.append(nn.BatchNorm1d(hidden_dim))
+            self.dropouts.append(nn.Dropout(dropout_rate))
+            prev_dim = hidden_dim
+
+        # Output layer
+        output_dim = num_points * point_dim
+        self.output_layer = nn.Linear(prev_dim, output_dim)
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        """
+        Decode latent code into point cloud.
+
+        Args:
+            z: Latent code of shape (B, latent_dim)
+
+        Returns:
+            Reconstructed point cloud of shape (B, num_points, point_dim)
+        """
+        x = z
+
+        # Process through hidden layers
+        for i, layer in enumerate(self.layers):
+            x = layer(x)
+            if self.use_batch_norm:
+                x = self.batch_norms[i](x)
+            x = self.activation(x)
+            x = self.dropouts[i](x)
+
+        # Output layer
+        x = self.output_layer(x)
+
+        # Reshape to point cloud
+        x = x.view(-1, self.num_points, self.point_dim)
+
+        return x
+
+    def get_output_shape(self) -> tuple:
+        """
+        Get the output point cloud shape (excluding batch dimension).
+
+        Returns:
+            Tuple of (num_points, point_dim)
+        """
+        return (self.num_points, self.point_dim)
+
+
+class ResidualMLPDecoder(nn.Module):
+    """
+    MLP decoder with residual connections for improved gradient flow.
+
+    This decoder uses residual blocks to help with training deep networks.
+
+    Args:
+        latent_dim: Dimensionality of input latent code
+        num_points: Number of output points in the point cloud
+        point_dim: Dimensionality of each point (default: 3 for XYZ coordinates)
+        hidden_dim: Hidden layer dimension (default: 1024)
+        num_blocks: Number of residual blocks (default: 3)
+        dropout_rate: Dropout probability (default: 0.3)
+        use_batch_norm: Whether to use batch normalization (default: True)
+
+    Input:
+        Latent code of shape (B, latent_dim)
+
+    Output:
+        Point cloud of shape (B, num_points, point_dim)
+    """
+
+    def __init__(
+        self,
+        latent_dim: int,
+        num_points: int,
+        point_dim: int = 3,
+        hidden_dim: int = 1024,
+        num_blocks: int = 3,
+        dropout_rate: float = 0.3,
+        use_batch_norm: bool = True
+    ):
+        super(ResidualMLPDecoder, self).__init__()
+        self.latent_dim = latent_dim
+        self.num_points = num_points
+        self.point_dim = point_dim
+        self.hidden_dim = hidden_dim
+
+        # Initial projection
+        self.input_proj = nn.Linear(latent_dim, hidden_dim)
+
+        # Residual blocks
+        self.blocks = nn.ModuleList([
+            ResidualBlock(hidden_dim, dropout_rate, use_batch_norm)
+            for _ in range(num_blocks)
+        ])
+
+        # Output projection
+        output_dim = num_points * point_dim
+        self.output_proj = nn.Linear(hidden_dim, output_dim)
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        """
+        Decode latent code into point cloud.
+
+        Args:
+            z: Latent code of shape (B, latent_dim)
+
+        Returns:
+            Reconstructed point cloud of shape (B, num_points, point_dim)
+        """
+        # Initial projection
+        x = F.relu(self.input_proj(z))
+
+        # Residual blocks
+        for block in self.blocks:
+            x = block(x)
+
+        # Output projection
+        x = self.output_proj(x)
+        x = x.view(-1, self.num_points, self.point_dim)
+
+        return x
+
+
+class ResidualBlock(nn.Module):
+    """
+    Residual block with optional batch normalization and dropout.
+
+    Args:
+        hidden_dim: Dimension of the hidden layer
+        dropout_rate: Dropout probability
+        use_batch_norm: Whether to use batch normalization
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        dropout_rate: float = 0.3,
+        use_batch_norm: bool = True
+    ):
+        super(ResidualBlock, self).__init__()
+        self.fc1 = nn.Linear(hidden_dim, hidden_dim)
+        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
+        self.dropout = nn.Dropout(dropout_rate)
+
+        if use_batch_norm:
+            self.bn1 = nn.BatchNorm1d(hidden_dim)
+            self.bn2 = nn.BatchNorm1d(hidden_dim)
+        else:
+            self.bn1 = None
+            self.bn2 = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass through residual block.
+
+        Args:
+            x: Input tensor of shape (B, hidden_dim)
+
+        Returns:
+            Output tensor of shape (B, hidden_dim)
+        """
+        residual = x
+
+        out = self.fc1(x)
+        if self.bn1 is not None:
+            out = self.bn1(out)
+        out = F.relu(out)
+        out = self.dropout(out)
+
+        out = self.fc2(out)
+        if self.bn2 is not None:
+            out = self.bn2(out)
+
+        out = out + residual
+        out = F.relu(out)
+
+        return out
+
+
+def create_pointcloud_decoder(
+    latent_dim: int,
+    num_points: int = 6890,
+    point_dim: int = 3,
+    architecture: str = 'mlp',
+    **kwargs
+) -> nn.Module:
+    """
+    Factory function to create a point cloud decoder with specified architecture.
+
+    Args:
+        latent_dim: Dimensionality of input latent code
+        num_points: Number of output points (default: 6890 for SMPL mesh)
+        point_dim: Dimensionality of each point (default: 3 for XYZ)
+        architecture: Decoder architecture type ('mlp' or 'residual')
+        **kwargs: Additional arguments passed to the decoder constructor
+
+    Returns:
+        Configured decoder instance
+
+    Examples:
+        >>> # Create a simple MLP decoder
+        >>> decoder = create_pointcloud_decoder(
+        ...     latent_dim=256,
+        ...     num_points=6890,
+        ...     architecture='mlp',
+        ...     hidden_dims=[1024, 2048],
+        ...     dropout_rate=0.3
+        ... )
+
+        >>> # Create a residual decoder
+        >>> decoder = create_pointcloud_decoder(
+        ...     latent_dim=256,
+        ...     num_points=6890,
+        ...     architecture='residual',
+        ...     hidden_dim=1024,
+        ...     num_blocks=3
+        ... )
+    """
+    if architecture == 'mlp':
+        return PointCloudDecoder(
+            latent_dim=latent_dim,
+            num_points=num_points,
+            point_dim=point_dim,
+            **kwargs
+        )
+    elif architecture == 'residual':
+        return ResidualMLPDecoder(
+            latent_dim=latent_dim,
+            num_points=num_points,
+            point_dim=point_dim,
+            **kwargs
+        )
+    else:
+        raise ValueError(f"Unsupported architecture: {architecture}")
+
+
+if __name__ == '__main__':
+    print("Testing Point Cloud Decoders...\n")
+
+    # Test parameters
+    batch_size = 32
+    latent_dim = 256
+    num_points = 6890  # SMPL mesh vertices
+    point_dim = 3
+
+    # Test standard MLP decoder
+    print("1. Standard MLP Decoder:")
+    mlp_decoder = create_pointcloud_decoder(
+        latent_dim=latent_dim,
+        num_points=num_points,
+        point_dim=point_dim,
+        architecture='mlp',
+        hidden_dims=[1024, 2048],
+        dropout_rate=0.3,
+        use_batch_norm=False
+    )
+
+    z = torch.randn(batch_size, latent_dim)
+    output = mlp_decoder(z)
+    print(f"  Input shape: {z.shape}")
+    print(f"  Output shape: {output.shape}")
+    print(f"  Output expected shape: {mlp_decoder.get_output_shape()}")
+    print()
+
+    # Test MLP decoder with batch norm
+    print("2. MLP Decoder with Batch Normalization:")
+    mlp_bn_decoder = create_pointcloud_decoder(
+        latent_dim=latent_dim,
+        num_points=num_points,
+        architecture='mlp',
+        use_batch_norm=True,
+        activation='leaky_relu'
+    )
+    output = mlp_bn_decoder(z)
+    print(f"  Output shape: {output.shape}")
+    print()
+
+    # Test residual MLP decoder
+    print("3. Residual MLP Decoder:")
+    residual_decoder = create_pointcloud_decoder(
+        latent_dim=latent_dim,
+        num_points=num_points,
+        architecture='residual',
+        hidden_dim=1024,
+        num_blocks=3,
+        dropout_rate=0.3
+    )
+    output = residual_decoder(z)
+    print(f"  Input shape: {z.shape}")
+    print(f"  Output shape: {output.shape}")
+    print()
+
+    # Test with different point dimensions
+    print("4. Decoder with 6D points (XYZ + RGB):")
+    decoder_6d = create_pointcloud_decoder(
+        latent_dim=latent_dim,
+        num_points=1000,
+        point_dim=6,
+        architecture='mlp',
+        hidden_dims=[512, 1024]
+    )
+    output = decoder_6d(z)
+    print(f"  Output shape: {output.shape}")
+    print()

src/generate_utils/lib/model/pointnet.py

@@ -0,0 +1,319 @@
+"""
+PointNet implementation for point cloud feature extraction.
+
+This module implements the PointNet architecture for encoding point clouds into
+global or per-point features. It can be used as a conditional encoder in cVAE models.
+
+Reference:
+    PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation
+    Charles R. Qi et al., CVPR 2017
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, Union
+
+
+class SpatialTransformer3D(nn.Module):
+    """
+    Spatial Transformer Network for 3D point clouds.
+
+    Predicts a 3x3 transformation matrix to canonicalize input point clouds.
+    """
+
+    def __init__(self):
+        super(SpatialTransformer3D, self).__init__()
+        self.conv1 = nn.Conv1d(3, 64, 1)
+        self.conv2 = nn.Conv1d(64, 128, 1)
+        self.conv3 = nn.Conv1d(128, 1024, 1)
+
+        self.fc1 = nn.Linear(1024, 512)
+        self.fc2 = nn.Linear(512, 256)
+        self.fc3 = nn.Linear(256, 9)
+
+        self.bn1 = nn.BatchNorm1d(64)
+        self.bn2 = nn.BatchNorm1d(128)
+        self.bn3 = nn.BatchNorm1d(1024)
+        self.bn4 = nn.BatchNorm1d(512)
+        self.bn5 = nn.BatchNorm1d(256)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the spatial transformer.
+
+        Args:
+            x: Input point cloud of shape (B, 3, N)
+
+        Returns:
+            Transformation matrix of shape (B, 3, 3)
+        """
+        batch_size = x.size(0)
+        device = x.device
+
+        # Encode point cloud
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = F.relu(self.bn2(self.conv2(x)))
+        x = F.relu(self.bn3(self.conv3(x)))
+
+        # Global max pooling
+        x = torch.max(x, 2, keepdim=True)[0]
+        x = x.view(batch_size, -1)
+
+        # Predict transformation
+        x = F.relu(self.bn4(self.fc1(x)))
+        x = F.relu(self.bn5(self.fc2(x)))
+        x = self.fc3(x)
+
+        # Add identity matrix as residual
+        identity = torch.eye(3, dtype=x.dtype, device=device).view(1, 9)
+        identity = identity.repeat(batch_size, 1)
+        x = x + identity
+        x = x.view(batch_size, 3, 3)
+
+        return x
+
+
+class SpatialTransformerKD(nn.Module):
+    """
+    Spatial Transformer Network for K-dimensional features.
+
+    Predicts a KxK transformation matrix for feature space alignment.
+
+    Args:
+        input_dim: Dimensionality of input features
+        feature_dim: Dimensionality of intermediate features (default: 1024)
+    """
+
+    def __init__(self, input_dim: int = 64, feature_dim: int = 1024):
+        super(SpatialTransformerKD, self).__init__()
+        self.input_dim = input_dim
+        self.feature_dim = feature_dim
+
+        self.conv1 = nn.Conv1d(input_dim, 64, 1)
+        self.conv2 = nn.Conv1d(64, 128, 1)
+        self.conv3 = nn.Conv1d(128, feature_dim, 1)
+
+        self.fc1 = nn.Linear(feature_dim, 512)
+        self.fc2 = nn.Linear(512, 256)
+        self.fc3 = nn.Linear(256, input_dim * input_dim)
+
+        self.bn1 = nn.BatchNorm1d(64)
+        self.bn2 = nn.BatchNorm1d(128)
+        self.bn3 = nn.BatchNorm1d(feature_dim)
+        self.bn4 = nn.BatchNorm1d(512)
+        self.bn5 = nn.BatchNorm1d(256)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the spatial transformer.
+
+        Args:
+            x: Input features of shape (B, K, N)
+
+        Returns:
+            Transformation matrix of shape (B, K, K)
+        """
+        batch_size = x.size(0)
+        device = x.device
+
+        # Encode features
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = F.relu(self.bn2(self.conv2(x)))
+        x = F.relu(self.bn3(self.conv3(x)))
+
+        # Global max pooling
+        x = torch.max(x, 2, keepdim=True)[0]
+        x = x.view(batch_size, -1)
+
+        # Predict transformation
+        x = F.relu(self.bn4(self.fc1(x)))
+        x = F.relu(self.bn5(self.fc2(x)))
+        x = self.fc3(x)
+
+        # Add identity matrix as residual
+        identity = torch.eye(self.input_dim, dtype=x.dtype, device=device)
+        identity = identity.view(1, self.input_dim * self.input_dim)
+        identity = identity.repeat(batch_size, 1)
+        x = x + identity
+        x = x.view(batch_size, self.input_dim, self.input_dim)
+
+        return x
+
+
+class PointNetEncoder(nn.Module):
+    """
+    PointNet encoder for extracting features from point clouds.
+
+    This encoder can output either global features (for the entire point cloud)
+    or per-point features (combining global and local information).
+
+    Args:
+        input_dim: Dimensionality of input point features (default: 3 for XYZ)
+        feature_dim: Dimensionality of output global features (default: 256)
+        use_spatial_transformer: Whether to use spatial transformer network (default: False)
+        return_global_feature: If True, return global feature; if False, return per-point features (default: True)
+
+    Input:
+        Point cloud of shape (B, input_dim, N) where:
+            B = batch size
+            input_dim = feature dimension per point (e.g., 3 for XYZ)
+            N = number of points
+
+    Output:
+        If return_global_feature=True:
+            Global feature of shape (B, feature_dim)
+        If return_global_feature=False:
+            Per-point features of shape (B, N, feature_dim + 64)
+    """
+
+    def __init__(
+        self,
+        input_dim: int = 3,
+        feature_dim: int = 256,
+        use_spatial_transformer: bool = False,
+        return_global_feature: bool = True
+    ):
+        super(PointNetEncoder, self).__init__()
+        self.input_dim = input_dim
+        self.feature_dim = feature_dim
+        self.use_spatial_transformer = use_spatial_transformer
+        self.return_global_feature = return_global_feature
+
+        # Spatial transformer
+        if use_spatial_transformer:
+            self.stn = SpatialTransformerKD(input_dim=input_dim, feature_dim=feature_dim)
+
+        # Feature extraction layers
+        self.conv1 = nn.Conv1d(input_dim, 64, 1)
+        self.conv2 = nn.Conv1d(64, 128, 1)
+        self.conv3 = nn.Conv1d(128, feature_dim, 1)
+
+        self.bn1 = nn.BatchNorm1d(64)
+        self.bn2 = nn.BatchNorm1d(128)
+        self.bn3 = nn.BatchNorm1d(feature_dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        return_transform: bool = False
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, Optional[torch.Tensor]]]:
+        """
+        Forward pass of PointNet encoder.
+
+        Args:
+            x: Input point cloud of shape (B, input_dim, N)
+            return_transform: If True, also return the transformation matrix (default: False)
+
+        Returns:
+            If return_transform=False:
+                features: Encoded features
+            If return_transform=True:
+                (features, transform): Tuple of features and transformation matrix
+        """
+        num_points = x.size(2)
+        transform = None
+
+        # Apply spatial transformation
+        if self.use_spatial_transformer:
+            transform = self.stn(x)
+            x = x.transpose(2, 1)  # (B, N, K)
+            x = torch.bmm(x, transform)  # (B, N, K)
+            x = x.transpose(2, 1)  # (B, K, N)
+
+        # Extract features
+        x = F.relu(self.bn1(self.conv1(x)))
+        local_features = x  # Save for per-point features
+
+        x = F.relu(self.bn2(self.conv2(x)))
+        x = self.bn3(self.conv3(x))
+
+        # Global max pooling
+        global_features = torch.max(x, 2, keepdim=True)[0]
+        global_features = global_features.view(-1, self.feature_dim)
+
+        # Return based on mode
+        if self.return_global_feature:
+            features = global_features
+        else:
+            # Concatenate global and local features for per-point features
+            global_features_expanded = global_features.view(-1, self.feature_dim, 1)
+            global_features_expanded = global_features_expanded.repeat(1, 1, num_points)
+            features = torch.cat([global_features_expanded, local_features], dim=1)
+            # Transpose to (B, N, feature_dim + 64)
+            features = features.transpose(1, 2)
+
+        if return_transform:
+            return features, transform
+        return features
+
+    def get_output_dim(self) -> int:
+        """
+        Get the output feature dimensionality.
+
+        Returns:
+            Output dimension based on the configuration
+        """
+        if self.return_global_feature:
+            return self.feature_dim
+        else:
+            return self.feature_dim + 64
+
+
+def create_pointnet_encoder(
+    input_dim: int = 3,
+    feature_dim: int = 256,
+    use_spatial_transformer: bool = False,
+    return_global_feature: bool = True
+) -> PointNetEncoder:
+    """
+    Factory function to create a PointNet encoder with specified configuration.
+
+    Args:
+        input_dim: Dimensionality of input point features
+        feature_dim: Dimensionality of output global features
+        use_spatial_transformer: Whether to use spatial transformer network
+        return_global_feature: Whether to return global or per-point features
+
+    Returns:
+        Configured PointNetEncoder instance
+    """
+    return PointNetEncoder(
+        input_dim=input_dim,
+        feature_dim=feature_dim,
+        use_spatial_transformer=use_spatial_transformer,
+        return_global_feature=return_global_feature
+    )
+
+
+if __name__ == '__main__':
+    print("Testing PointNetEncoder with different configurations...\n")
+
+    # Test data
+    batch_size = 32
+    num_points = 6890  # SMPL mesh vertices
+    input_dim = 3
+
+    sim_data = torch.randn(batch_size, input_dim, num_points)
+
+    configs = [
+        (True, False, "Global features without STN"),
+        (True, True, "Global features with STN"),
+        (False, False, "Per-point features without STN"),
+        (False, True, "Per-point features with STN"),
+    ]
+
+    for return_global, use_stn, desc in configs:
+        encoder = create_pointnet_encoder(
+            input_dim=input_dim,
+            feature_dim=256,
+            use_spatial_transformer=use_stn,
+            return_global_feature=return_global
+        )
+
+        output = encoder(sim_data)
+        print(f"{desc}:")
+        print(f"  Input shape: {sim_data.shape}")
+        print(f"  Output shape: {output.shape}")
+        print(f"  Output dim: {encoder.get_output_dim()}")
+        print()

src/generate_utils/lib/model/resnet.py

@@ -0,0 +1,432 @@
+import torch
+import torch.nn as nn
+from torch.hub import load_state_dict_from_url
+
+
+__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
+           'resnet152', 'resnext50_32x4d', 'resnext101_32x8d',
+           'wide_resnet50_2', 'wide_resnet101_2']
+
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+    'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
+    'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
+    'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
+    'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+    # This variant is also known as ResNet V1.5 and improves accuracy according to
+    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(Bottleneck, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.)) * groups
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.conv2 = conv3x3(width, width, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+        self.conv3 = conv1x1(width, planes * self.expansion)
+        self.bn3 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+    """
+    ResNet encoder for conditional VAE.
+
+    Args:
+        block: Type of residual block (BasicBlock or Bottleneck)
+        layers: List of integers indicating the number of blocks in each layer
+        embed_dim: Dimension of the latent embedding (output dimension)
+        zero_init_residual: Whether to zero-initialize the last BN in each residual branch
+        groups: Number of groups for grouped convolutions
+        width_per_group: Width per group for grouped convolutions
+        replace_stride_with_dilation: List of booleans for using dilated convolutions
+        norm_layer: Normalization layer to use (default: BatchNorm2d)
+        cond_embed_dim: Dimension of the conditional embedding input
+        dp_rate: Dropout rate applied before the final linear layers
+
+    Inputs:
+        x: Input tensor of shape (batch_size, 1, height, width)
+        cond: Conditional embedding of shape (batch_size, cond_embed_dim)
+
+    Outputs:
+        mu: Mean of the latent distribution, shape (batch_size, embed_dim)
+        log_var: Log variance of the latent distribution, shape (batch_size, embed_dim)
+    """
+    def __init__(self, block, layers, embed_dim=256, zero_init_residual=False,
+                 groups=1, width_per_group=64, replace_stride_with_dilation=None,
+                 norm_layer=None, cond_embed_dim=256, dp_rate=0.3):
+        super(ResNet, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            # each element in the tuple indicates if we should replace
+            # the 2x2 stride with a dilated convolution instead
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(1, self.inplanes, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
+                                       dilate=replace_stride_with_dilation[0])
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
+                                       dilate=replace_stride_with_dilation[1])
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
+                                       dilate=replace_stride_with_dilation[2])
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+
+        # Condition embedding network to align with visual features
+        self.cond_fc = nn.Sequential(
+            nn.Linear(cond_embed_dim, 512*block.expansion),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dp_rate)
+        )
+
+        # Shared fusion layer before mu and log_var prediction
+        fusion_dim = 512 * block.expansion
+        self.fusion_fc = nn.Sequential(
+            nn.Linear(512*block.expansion*2, fusion_dim),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dp_rate)
+        )
+
+        # Separate heads for mu and log_var with independent dropout
+        self.fc_mu = nn.Sequential(
+            nn.Dropout(dp_rate),
+            nn.Linear(fusion_dim, embed_dim)
+        )
+        self.fc_logvar = nn.Sequential(
+            nn.Dropout(dp_rate),
+            nn.Linear(fusion_dim, embed_dim)
+        )
+        
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
+                            self.base_width, previous_dilation, norm_layer))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, groups=self.groups,
+                                base_width=self.base_width, dilation=self.dilation,
+                                norm_layer=norm_layer))
+
+        return nn.Sequential(*layers)
+
+    def _forward_impl(self, x, cond):
+        # Extract visual features
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+
+        # Process condition through embedding network
+        cond_embed = self.cond_fc(cond)
+
+        # Concatenate visual features and condition embedding
+        x = torch.cat((x, cond_embed), dim=1)
+
+        # Shared fusion
+        x = self.fusion_fc(x)
+
+        # Predict mu and log_var with independent dropout
+        mu = self.fc_mu(x)
+        log_var = self.fc_logvar(x)
+
+        return mu, log_var
+
+    def forward(self, x, cond):
+        return self._forward_impl(x, cond)
+
+
+def _resnet(arch, block, layers, pretrained, progress, **kwargs):
+    model = ResNet(block, layers, **kwargs)
+    if pretrained:
+        state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
+        # Handle the mismatch in conv1 (3 channels -> 1 channel)
+        if 'conv1.weight' in state_dict:
+            pretrained_conv1 = state_dict['conv1.weight']
+            # Average the weights across RGB channels for single channel input
+            state_dict['conv1.weight'] = pretrained_conv1.mean(dim=1, keepdim=True)
+
+        # Remove classifier weights as we have custom fc layers
+        state_dict = {k: v for k, v in state_dict.items() if not k.startswith('fc')}
+        model.load_state_dict(state_dict, strict=False)
+    return model
+
+
+def resnet18(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-18 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress, **kwargs)
+
+
+def resnet34(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-34 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress, **kwargs)
+
+
+def resnet50(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-50 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs)
+
+
+def resnet101(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-101 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs)
+
+
+def resnet152(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-152 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress, **kwargs)
+
+
+def resnext50_32x4d(pretrained=False, progress=True, **kwargs):
+    r"""ResNeXt-50 32x4d model from
+    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['groups'] = 32
+    kwargs['width_per_group'] = 4
+    return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3],
+                   pretrained, progress, **kwargs)
+
+
+def resnext101_32x8d(pretrained=False, progress=True, **kwargs):
+    r"""ResNeXt-101 32x8d model from
+    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['groups'] = 32
+    kwargs['width_per_group'] = 8
+    return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3],
+                   pretrained, progress, **kwargs)
+
+
+def wide_resnet50_2(pretrained=False, progress=True, **kwargs):
+    r"""Wide ResNet-50-2 model from
+    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_
+
+    The model is the same as ResNet except for the bottleneck number of channels
+    which is twice larger in every block. The number of channels in outer 1x1
+    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
+    channels, and in Wide ResNet-50-2 has 2048-1024-2048.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['width_per_group'] = 64 * 2
+    return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3],
+                   pretrained, progress, **kwargs)
+
+
+def wide_resnet101_2(pretrained=False, progress=True, **kwargs):
+    r"""Wide ResNet-101-2 model from
+    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_
+
+    The model is the same as ResNet except for the bottleneck number of channels
+    which is twice larger in every block. The number of channels in outer 1x1
+    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
+    channels, and in Wide ResNet-50-2 has 2048-1024-2048.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['width_per_group'] = 64 * 2
+    return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3],
+                   pretrained, progress, **kwargs)
+
+
+if __name__ == "__main__":
+    a = torch.randn(8, 1, 110, 37)
+    cond = torch.randn(8, 85)
+    encoder = resnet18(embed_dim=256, cond_embed_dim=85)
+    b, var = encoder(a, cond)
+    print(b.shape)
+    print(var.shape)
+
+    encoder = resnet34(embed_dim=256, cond_embed_dim=85)
+    b, var = encoder(a, cond)
+    print(b.shape)
+    print(var.shape)
+    
+    encoder = resnet50(embed_dim=1024, cond_embed_dim=85)
+    b, var = encoder(a, cond)
+    print(b.shape)
+    print(var.shape)

src/generate_utils/lib/model/unet.py

@@ -0,0 +1,277 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def get_crop_config(crop):
+    """
+    Get configuration parameters based on crop size.
+
+    Args:
+        crop: List of [height, width] or None
+
+    Returns:
+        dict: Configuration including alpha, h, w, diffx, diffy
+    """
+    if crop is None:
+        # Default values for TIP dataset (56x40)
+        return {
+            'alpha': 6,
+            'h': 3,
+            'w': 2,
+            'diffx': [1, 0, 0, 0],
+            'diffy': [1, 0, 0, 0]
+        }
+    elif crop == [64, 27]:  # PressurePose dataset
+        return {
+            'alpha': 4,
+            'h': 4,
+            'w': 1,
+            'diffx': [1, 0, 1, 1],
+            'diffy': [0, 0, 0, 0]
+        }
+    elif crop == [110, 37]:  # MOYO dataset
+        return {
+            'alpha': 12,
+            'h': 6,
+            'w': 2,
+            'diffx': [0, 1, 0, 1],
+            'diffy': [1, 1, 1, 0]
+        }
+    else:  # Default for other sizes
+        return {
+            'alpha': 6,
+            'h': 3,
+            'w': 2,
+            'diffx': [1, 0, 0, 0],
+            'diffy': [1, 0, 0, 0]
+        }
+
+
+class DoubleConv(nn.Module):
+    """Double Convolutional Block with BN and ReLU"""
+    
+    def __init__(self, in_channels, out_channels, mid_channels=None):
+        super().__init__()
+        if not mid_channels:
+            mid_channels = out_channels
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+    
+    def forward(self, x):
+        return self.double_conv(x)
+
+
+class Down(nn.Module):
+    """Downscaling with maxpool then double conv"""
+    
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            DoubleConv(in_channels, out_channels)
+        )
+    
+    def forward(self, x):
+        return self.maxpool_conv(x)
+
+
+class Up(nn.Module):
+    """Upscaling then double conv"""
+    
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
+        else:
+            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
+            self.conv = DoubleConv(in_channels, out_channels)
+    
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        diffY = x2.size()[2] - x1.size()[2]
+        diffX = x2.size()[3] - x1.size()[3]
+
+        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2])
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+
+class OutConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(OutConv, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+    
+    def forward(self, x):
+        return self.conv(x)
+
+
+class ConditionalBatchNorm2d(nn.Module):
+    """Conditional BatchNorm2d to modulate output with condition vector"""
+    
+    def __init__(self, num_features, cond_dim):
+        super(ConditionalBatchNorm2d, self).__init__()
+        self.num_features = num_features
+        self.bn = nn.BatchNorm2d(num_features, affine=False)
+        self.gamma = nn.Linear(cond_dim, num_features)
+        self.beta = nn.Linear(cond_dim, num_features)
+    
+    def forward(self, x, cond):
+        gamma = self.gamma(cond).view(-1, self.num_features, 1, 1)
+        beta = self.beta(cond).view(-1, self.num_features, 1, 1)
+        return self.bn(x) * gamma + beta
+
+
+class UpWithCondition(nn.Module):
+    """Upscaling then double conv with condition modulation"""
+    
+    def __init__(self, in_channels, out_channels, cond_dim, diffX, diffY, bilinear=True):
+        super(UpWithCondition, self).__init__()
+        self.diffX = diffX
+        self.diffY = diffY
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
+        else:
+            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
+            self.conv = DoubleConv(in_channels // 2, out_channels)
+
+        # Conditional batch norm for output modulation
+        self.cond_bn = ConditionalBatchNorm2d(out_channels, cond_dim)
+    
+    def forward(self, x, cond):
+        x = self.up(x)
+        x = F.pad(x, [self.diffX // 2, self.diffX - self.diffX // 2, self.diffY // 2, self.diffY - self.diffY // 2])
+        x = self.conv(x)
+        x = self.cond_bn(x, cond)  # Modulate with condition
+        return x
+
+
+class UNetEncoder(nn.Module):
+    def __init__(self, cond_dim=256, embed_dim=256, dp_rate=0.0, bilinear=False, crop=None):
+        super(UNetEncoder, self).__init__()
+        self.cond_dim = cond_dim
+        self.crop = crop
+
+        # Get crop-specific configuration
+        crop_config = get_crop_config(crop)
+        alpha = crop_config['alpha']
+
+        # Encoder
+        self.inc = DoubleConv(1, 64)
+        self.down1 = Down(64, 128)
+        self.down2 = Down(128, 256)
+        self.down3 = Down(256, 512)
+        factor = 2 if bilinear else 1
+        self.down4 = Down(512, 1024 // factor)
+
+        self.dropout = nn.Dropout(dp_rate)
+
+        # VAE latent space parameters
+        self.fc_mu = nn.Linear((1024 // factor) * alpha + cond_dim, embed_dim)
+        self.fc_log_var = nn.Linear((1024 // factor) * alpha + cond_dim, embed_dim)
+    
+    def forward(self, x, cond):
+        # Encoder path
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)  # B x (1024 // factor) x H x W
+
+        # Flatten and concatenate with condition vector
+        x5_flat = x5.view(x5.size(0), -1)
+        x5_cond = torch.cat([x5_flat, cond], dim=1)
+        
+        # Compute mu and log_var for latent space
+        mu = self.fc_mu(x5_cond)
+        log_var = self.fc_log_var(x5_cond)
+        
+        return mu, log_var
+        
+
+class UNetDecoder(nn.Module):
+    def __init__(self, cond_dim=256, embed_dim=256, dp_rate=0.0, bilinear=False, crop=None):
+        super(UNetDecoder, self).__init__()
+        self.bilinear = bilinear
+        self.cond_dim = cond_dim
+        self.crop = crop
+
+        # Get crop-specific configuration
+        crop_config = get_crop_config(crop)
+        alpha = crop_config['alpha']
+        self.h = crop_config['h']
+        self.w = crop_config['w']
+        diffx = crop_config['diffx']
+        diffy = crop_config['diffy']
+
+        factor = 2 if bilinear else 1
+
+        # Map latent vector and condition back to decoder size
+        self.fc_z = nn.Linear(embed_dim + cond_dim, (1024 // factor) * alpha)
+        self.dropout = nn.Dropout(dp_rate)
+
+        # Decoder with conditional batch norm
+        self.up1 = UpWithCondition(1024, 512 // factor, cond_dim, diffX=diffx[0], diffY=diffy[0], bilinear=bilinear)
+        self.up2 = UpWithCondition(512, 256 // factor, cond_dim, diffX=diffx[1], diffY=diffy[1], bilinear=bilinear)
+        self.up3 = UpWithCondition(256, 128 // factor, cond_dim, diffX=diffx[2], diffY=diffy[2], bilinear=bilinear)
+        self.up4 = UpWithCondition(128, 64, cond_dim, diffX=diffx[3], diffY=diffy[3], bilinear=bilinear)
+        self.outc = OutConv(64, 1)
+    
+    def forward(self, z, cond):
+        b, _ = z.shape
+        factor = 2 if self.bilinear else 1
+
+        # Decode latent vector concatenated with condition
+        z_cond = torch.cat([z, cond], dim=1)
+        z_decoded = self.fc_z(z_cond).view(b, 1024 // factor, self.h, self.w)
+        z_decoded = self.dropout(z_decoded)
+
+        # Decoder path with conditional modulation
+        x = self.up1(z_decoded, cond)
+        x = self.up2(x, cond)
+        x = self.up3(x, cond)
+        x = self.up4(x, cond)
+        x = self.outc(x)
+
+        return x.view(b, self.crop[0], self.crop[1])
+
+
+if __name__ == "__main__":
+    def reparameterize(mu, logvar):
+        """VAE reparameterization trick"""
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        return mu + eps * std
+
+    # Test data
+    data = torch.randn(4, 1, 56, 40)
+    condition = torch.randn(4, 85)
+
+    # Initialize encoder and decoder
+    encoder = UNetEncoder(cond_dim=85, embed_dim=256, bilinear=False, crop=[56, 40])
+    decoder = UNetDecoder(cond_dim=85, embed_dim=256, bilinear=False, crop=[56, 40])
+
+    # Forward pass through encoder
+    mu, logvar = encoder(data, condition)
+    print(f"Mu shape: {mu.shape}, Logvar shape: {logvar.shape}")
+
+    # Reparameterization trick
+    z = reparameterize(mu, logvar)
+    print(f"Latent z shape: {z.shape}")
+
+    # Forward pass through decoder
+    reconstructed = decoder(z, condition)
+    print(f"Reconstructed shape: {reconstructed.shape}")
+    print(f"Expected shape: [4, 56, 40]")
+
+    assert reconstructed.shape == (4, 56, 40), f"Shape mismatch: {reconstructed.shape} vs (4, 56, 40)"
+    print("Test passed!")
+