Refactored code for AddBiomechanics and previous datasets

Author: Starfractor

Dockerfile

@@ -1,11 +1,11 @@
 # Use an official Ubuntu runtime as a parent image
-FROM ubuntu:22.04
+FROM ubuntu:25.04
 
 # Set the working directory
-WORKDIR /P-BIGE
+WORKDIR /BIGE
 
 # Install necessary dependencies
-RUN apt-get update && apt-get install -y wget git htop xvfb
+RUN apt-get update && apt-get install -y wget git htop xvfb build-essential 
 
 # Install Miniconda
 RUN MINICONDA_INSTALLER_SCRIPT=Miniconda3-py38_23.1.0-1-Linux-x86_64.sh && \
@@ -19,25 +19,21 @@ RUN MINICONDA_INSTALLER_SCRIPT=Miniconda3-py38_23.1.0-1-Linux-x86_64.sh && \
 ENV PATH=/usr/local/bin:$PATH
 
 # Copy the entire repo (including environment.yml) into the image
-COPY . /P-BIGE
+COPY . /BIGE
 
 # Create the conda environment from environment.yml
-RUN conda env create -f environment.yml
+RUN conda env create -f enviornment.yaml
 
 # Activate the conda environment for subsequent RUN commands
-SHELL ["conda", "run", "-n", "P-BIGE", "/bin/bash", "-c"]
+SHELL ["conda", "run", "-n", "BIGE", "/bin/bash", "-c"]
 
-# Download the model and extractor
-RUN bash dataset/prepare/download_model.sh && \
-    bash dataset/prepare/download_extractor.sh
+# # Install additional Python packages (if needed)
+# RUN pip install --user ipykernel polyscope easydict trimesh
+# RUN pip install --user --force-reinstall numpy==1.22.0
 
-# Install additional Python packages (if needed)
-RUN pip install --user ipykernel polyscope easydict trimesh
-RUN pip install --user --force-reinstall numpy==1.22.0
-
-# Set up Xvfb for Polyscope
-ENV DISPLAY=:99.0
-RUN Xvfb :99 -screen 0 1024x768x24 > /dev/null 2>&1 &
+# # Set up Xvfb for Polyscope
+# ENV DISPLAY=:99.0
+# RUN Xvfb :99 -screen 0 1024x768x24 > /dev/null 2>&1 &
 
 # Set the entrypoint to use the conda environment
-ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "P-BIGE", "python"]
+ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "BIGE", "python"]

dataset/dataset_MOT_segmented.py

@@ -205,7 +205,7 @@ 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'):
+    def __init__(self, window_size=64, unit_length=4, mode='train', data_dir='/home/kingn450/Datasets/addb_dataset_publication'):
         self.window_size = window_size
         self.unit_length = unit_length
         self.data_dir = data_dir

dev-pod.yaml

@@ -0,0 +1,23 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  name: dev-pod
+spec:
+  containers:
+  - name: sleep
+    image: ncking/p-bige:latest
+    command: ["/bin/bash", "-c"]
+    args: ["sleep infinity"]
+    resources:
+      limits:
+        nvidia.com/gpu: 1
+      requests:
+        nvidia.com/gpu: 1
+    volumeMounts:
+    - name: biomechanics-dataset
+      mountPath: /home/mnt/data
+  volumes:
+  - name: biomechanics-dataset
+    persistentVolumeClaim:
+      claimName: add-biomechanics-dataset-pvc
+  restartPolicy: Never

enviornment.yaml

@@ -0,0 +1,17 @@
+name: BIGE
+channels:
+  - defaults
+  - nvidia
+  - pytorch
+dependencies:
+  - cudatoolkit
+  - imageio
+  - matplotlib
+  - python=3.10
+  - scipy
+  - tensorboard
+  - pip:
+    - deepspeed
+    - git+https://github.com/nghorbani/human_body_prior
+    - git+https://github.com/openai/CLIP.git
+    - nimblephysics

environment.yml

@@ -1,4 +1,4 @@
-name: P-BIGE
+name: T2M-GPT
 channels:
   - pytorch
   - defaults
@@ -71,26 +71,20 @@ dependencies:
     - chumpy==0.70
     - cycler==0.10.0
     - decorator==5.0.9
-    - deepspeed==0.5.8
-    - gdown
-    - git+https://github.com/nghorbani/human_body_prior
-    - git+https://github.com/openai/CLIP.git
-    - google-auth-oauthlib==0.4.5
     - google-auth==1.35.0
+    - google-auth-oauthlib==0.4.5
     - grpcio==1.39.0
     - idna==3.2
     - imageio==2.9.0
     - ipdb==0.13.9
-    - ipython-genutils==0.2.0
     - ipython==7.26.0
+    - ipython-genutils==0.2.0
     - jedi==0.18.0
     - joblib==1.0.1
     - kiwisolver==1.3.1
     - markdown==3.3.4
-    - matplotlib-inline==0.1.2
     - matplotlib==3.4.3
-    - moviepy
-    - nimblephysics
+    - matplotlib-inline==0.1.2
     - oauthlib==3.1.1
     - pandas==1.3.2
     - parso==0.8.2
@@ -99,27 +93,31 @@ dependencies:
     - prompt-toolkit==3.0.20
     - protobuf==3.17.3
     - ptyprocess==0.7.0
-    - pyasn1-modules==0.2.8
     - pyasn1==0.4.8
+    - pyasn1-modules==0.2.8
     - pygments==2.10.0
     - pyparsing==2.4.7
     - python-dateutil==2.8.2
     - pytz==2021.1
     - pyyaml==5.4.1
-    - requests-oauthlib==1.3.0
     - requests==2.26.0
+    - requests-oauthlib==1.3.0
     - rsa==4.7.2
     - scikit-learn==0.24.2
     - scipy==1.7.1
     - sklearn==0.0
     - smplx==0.1.28
+    - tensorboard==2.6.0
     - tensorboard-data-server==0.6.1
     - tensorboard-plugin-wit==1.8.0
-    - tensorboard==2.6.0
     - threadpoolctl==2.2.0
     - toml==0.10.2
     - tqdm==4.62.2
     - traitlets==5.0.5
     - urllib3==1.26.6
     - wcwidth==0.2.5
-    - werkzeug==2.0.1
+    - werkzeug==2.0.1
+    - git+https://github.com/openai/CLIP.git
+    - git+https://github.com/nghorbani/human_body_prior
+    - gdown
+    - moviepy

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(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)
+        self.encoder = Encoder(23 if args.dataname == 'mcs' or args.dataname == 'addb' 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' or args.dataname == 'addb' 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":