Updated dataloader to work with AddBiomechanics dataset

Author: Starfractor

Dockerfile

@@ -22,15 +22,16 @@ ENV PATH=/usr/local/bin:$PATH
 
 # Clone UCSD-Github dataset 
 # Set the working directory
-WORKDIR /
-RUN git -c http.sslVerify=false clone https://github.com/Rose-STL-Lab/UCSD-OpenCap-Fitness-Dataset.git
+#WORKDIR /
+#RUN git -c http.sslVerify=false clone https://github.com/Rose-STL-Lab/UCSD-OpenCap-Fitness-Dataset.git
 
 
 # Clone the digital-coach-anwesh repository
-RUN git -c http.sslVerify=false clone https://gitlab.nrp-nautilus.io/shmaheshwari/digital-coach-anwesh.git .
+#RUN git -c http.sslVerify=false clone https://gitlab.nrp-nautilus.io/shmaheshwari/digital-coach-anwesh.git .
 
 # Copy the environment.yml file and create the conda environment
 # COPY digital-coach-anwesh/environment.yml /T2M-GPT/environment.yml
+COPY . /T2M-GPT
 RUN conda env create -f environment.yml
 
 # Activate the conda environment

dataset/dataset_MOT_segmented.py

@@ -1,13 +1,13 @@
+import codecs as cs
+import nimblephysics as nimble
+import numpy as np
 import os
+import random
 import torch
-from torch.utils import data
-import numpy as np
+from glob import glob
 from os.path import join as pjoin
-import random
-import codecs as cs
+from torch.utils import data
 from tqdm import tqdm
-from glob import glob
-
 
 class VQMotionDataset(data.Dataset):
     def __init__(self, dataset_name, window_size = 64, unit_length = 4, mode = 'train', mode2='embeddings', data_dirs=['/home/ubuntu/data/MCS_DATA', '/media/shubh/Elements/RoseYu/UCSD-OpenCap-Fitness-Dataset/MCS_DATA']):
@@ -204,6 +204,108 @@ def __getitem__(self, item):
 
             return subsequences, subsequence_lengths, names
 
+class AddBiomechanicsDataset(data.Dataset):
+    def __init__(self, window_size=64, unit_length=4, mode='train', data_dir='/home/mnt/data/addb_dataset_publication'):
+        self.window_size = window_size
+        self.unit_length = unit_length
+        self.data_dir = data_dir
+        self.mode = mode
+
+        # Define subdirectories for each paper
+        paper_dirs = [
+            "train/No_Arm/Falisse2016_Formatted_No_Arm",
+            "train/No_Arm/Uhlrich2023_Opencap_Formatted_No_Arm",
+            "train/No_Arm/Wang2023_Formatted_No_Arm",
+            "train/No_Arm/Han2023_Formatted_No_Arm",
+        ]
+
+        # Collect all .b3d files from the specified subdirectories
+        self.b3d_file_paths = []
+        for paper_dir in paper_dirs:
+            search_path = os.path.join(data_dir, paper_dir, '**', '*.b3d')
+            files = glob(search_path, recursive=True)
+            self.b3d_file_paths.extend(files)
+
+        self.motion_data = []
+        self.motion_lengths = []
+        self.motion_names = []
+        self.motion_fps = []
+
+        for b3d_file in tqdm(self.b3d_file_paths):
+            try:
+                if os.path.getsize(b3d_file) == 0:
+                    continue
+                subject = nimble.biomechanics.SubjectOnDisk(b3d_file)
+                num_trials = subject.getNumTrials()
+                for trial in range(num_trials):
+                    trial_length = subject.getTrialLength(trial)
+                    if trial_length < self.window_size:
+                        continue
+                    frames = subject.readFrames(
+                        trial=trial,
+                        startFrame=0,
+                        numFramesToRead=trial_length,
+                        includeSensorData=False,
+                        includeProcessingPasses=True
+                    )
+                    if not frames:
+                        continue
+                    kin_passes = [frame.processingPasses[0] for frame in frames]
+                    positions = np.array([kp.pos for kp in kin_passes])  # shape: (frames, dofs)
+                    # Get FPS for this trial
+                    seconds_per_frame = subject.getTrialTimestep(trial)
+                    fps = int(round(1.0 / seconds_per_frame)) if seconds_per_frame > 0 else 0
+
+                    # Downsample here, at load time
+                    if fps == 100:
+                        positions = positions[::2]  # Take every 2nd frame
+                    elif fps == 250:
+                        positions = positions[::5]  # Take every 5th frame
+
+                    # After downsampling, skip if too short
+                    if len(positions) < self.window_size:
+                        continue
+
+                    self.motion_data.append(positions)
+                    self.motion_lengths.append(len(positions))
+                    self.motion_names.append(f"{b3d_file}::trial{trial}")
+                    self.motion_fps.append(fps)
+            except Exception as e:
+                print(f"Skipping file {b3d_file} due to error: {e}")
+
+        print("Total number of motions:", len(self.motion_data))
+        print("Example motion shape:", self.motion_data[0].shape if self.motion_data else "None")
+
+    def __len__(self):
+        return len(self.motion_data)
+
+    def __getitem__(self, item):
+        motion = self.motion_data[item]
+        len_motion = len(motion) if len(motion) <= self.window_size else self.window_size
+        name = self.motion_names[item]
+
+        # Crop or pad to window_size (no downsampling here)
+        if len(motion) >= self.window_size:
+            idx = random.randint(0, len(motion) - self.window_size)
+            motion = motion[idx:idx + self.window_size]
+        else:
+            repeat_count = (self.window_size + len(motion) - 1) // len(motion)
+            motion = np.tile(motion, (repeat_count, 1))[:self.window_size]
+
+        return motion, len_motion, name
+    
+
+def addb_data_loader(window_size=64, unit_length=4, batch_size=1, num_workers=4, mode='train'):
+    dataset = AddBiomechanicsDataset(window_size=window_size, unit_length=unit_length, mode=mode)
+    loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        drop_last=True
+    )
+    return loader
+
 def DATALoader(dataset_name,
                batch_size,
                num_workers = 4,
@@ -231,4 +333,4 @@ def cycle(iterable):
 
 
 if __name__ == "__main__": 
-    dataloader = DATALoader('mcs',1,window_size=64,unit_length=2**2,mode='limo')
+    dataloader = addb_data_loader(window_size=64, unit_length=4, batch_size=1, mode='train')

environment.yml

@@ -85,6 +85,7 @@ dependencies:
     - markdown==3.3.4
     - matplotlib==3.4.3
     - matplotlib-inline==0.1.2
+    - nimblephysics
     - oauthlib==3.1.1
     - pandas==1.3.2
     - parso==0.8.2

models/vqvae.py

@@ -21,8 +21,8 @@ def __init__(self,
         self.code_dim = code_dim
         self.num_code = nb_code
         self.quant = args.quantizer
-        self.encoder = Encoder(33 if args.dataname == 'mcs' else 263, output_emb_width, down_t, stride_t, width, depth, dilation_growth_rate, activation=activation, norm=norm)
-        self.decoder = Decoder(33 if args.dataname == 'mcs' else 263, output_emb_width, down_t, stride_t, width, depth, dilation_growth_rate, activation=activation, norm=norm)
+        self.encoder = Encoder(23 if args.dataname == 'mcs' else 263, output_emb_width, down_t, stride_t, width, depth, dilation_growth_rate, activation=activation, norm=norm)
+        self.decoder = Decoder(23 if args.dataname == 'mcs' else 263, output_emb_width, down_t, stride_t, width, depth, dilation_growth_rate, activation=activation, norm=norm)
         if args.quantizer == "ema_reset":
             self.quantizer = QuantizeEMAReset(nb_code, code_dim, args)
         elif args.quantizer == "orig":

train_vq.py

@@ -17,7 +17,7 @@
 import warnings
 warnings.filterwarnings('ignore')
 from utils.word_vectorizer import WordVectorizer
-# import nimblephysics as nimble
+import nimblephysics as nimble
 import deepspeed
 
 
@@ -67,7 +67,8 @@ def get_foot_losses(motion, y_translation=0.0,feet_threshold=0.01):
 ##### ---- Exp dirs ---- #####
 args = option_vq.get_args_parser()
 torch.manual_seed(args.seed)
-torch.cuda.set_device(args.local_rank)
+if torch.cuda.is_available():
+    torch.cuda.set_device(args.local_rank)
 
 args.out_dir = os.path.join(args.out_dir, f'{args.exp_name}')
 os.makedirs(args.out_dir, exist_ok = True)
@@ -87,10 +88,13 @@ def get_foot_losses(motion, y_translation=0.0,feet_threshold=0.01):
     dataset_opt_path = 'checkpoints/t2m/Comp_v6_KLD005/opt.txt'
     args.nb_joints = 22
 
+args.nb_joints = 23 # fixed issues
+
 logger.info(f'Training on {args.dataname}, motions are with {args.nb_joints} joints')
 
-wrapper_opt = get_opt(dataset_opt_path, torch.device('cuda'))
-eval_wrapper = EvaluatorModelWrapper(wrapper_opt)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+wrapper_opt = get_opt(dataset_opt_path, device)
+#eval_wrapper = EvaluatorModelWrapper(wrapper_opt)
 
 
 ##### ---- Dataloader ---- #####
@@ -99,18 +103,20 @@ def get_foot_losses(motion, y_translation=0.0,feet_threshold=0.01):
 #                                         window_size=args.window_size,
 #                                         unit_length=2**args.down_t)
 
-train_loader = dataset_MOT_segmented.DATALoader(args.dataname,
-                                        args.batch_size,
-                                        window_size=args.window_size,
-                                        unit_length=2**args.down_t)
+train_loader = dataset_MOT_segmented.addb_data_loader(
+    window_size=args.window_size,
+    unit_length=2**args.down_t,
+    batch_size=args.batch_size,
+    mode=args.dataname
+)
 
 # train_loader_iter = dataset_MOT_MCS.cycle(train_loader)
 train_loader_iter = dataset_MOT_segmented.cycle(train_loader)
 
-val_loader = dataset_TM_eval.DATALoader(args.dataname, False,
-                                        32,
-                                        w_vectorizer,
-                                        unit_length=2**args.down_t)
+# val_loader = dataset_TM_eval.DATALoader(args.dataname, False,
+#                                         32,
+#                                         w_vectorizer,
+#                                         unit_length=2**args.down_t)
 
 ##### ---- Network ---- #####
 net = vqvae.HumanVQVAE(args, ## use args to define different parameters in different quantizers
@@ -128,10 +134,10 @@ def get_foot_losses(motion, y_translation=0.0,feet_threshold=0.01):
 
 if args.resume_pth : 
     logger.info('loading checkpoint from {}'.format(args.resume_pth))
-    ckpt = torch.load(args.resume_pth, map_location='cuda')
+    ckpt = torch.load(args.resume_pth, map_location=device)
     net.load_state_dict(ckpt['net'], strict=True)
 net.train()
-net.cuda()
+net.to(device)
 
 ##### ---- Optimizer & Scheduler ---- #####
 optimizer = optim.AdamW(net.parameters(), lr=args.lr, betas=(0.9, 0.99), weight_decay=args.weight_decay)
@@ -170,7 +176,7 @@ def get_foot_losses(motion, y_translation=0.0,feet_threshold=0.01):
     optimizer, current_lr = update_lr_warm_up(optimizer, nb_iter, args.warm_up_iter, args.lr)
 
     gt_motion,_, names = next(train_loader_iter)
-    gt_motion = gt_motion.cuda().float() # (bs, 64, dim)
+    gt_motion = gt_motion.to(device).float() # (bs, 64, dim)
 
     pred_motion, loss_commit, perplexity = net(gt_motion)
     loss_motion = Loss(pred_motion, gt_motion)
@@ -222,7 +228,7 @@ def get_foot_losses(motion, y_translation=0.0,feet_threshold=0.01):
 for nb_iter in range(1, args.total_iter + 1):
 
     gt_motion,_,_ = next(train_loader_iter)
-    gt_motion = gt_motion.cuda().float() # bs, nb_joints, joints_dim, seq_len
+    gt_motion = gt_motion.to(device).float() # bs, nb_joints, joints_dim, seq_len
 
     pred_motion, loss_commit, perplexity = net(gt_motion)
     loss_motion = Loss(pred_motion, gt_motion)