Predicting 6D state for each body segment

.gitignore

@@ -1,5 +1,24 @@
-# Created by https://www.toptal.com/developers/gitignore/api/python,linux,macos,windows,jupyternotebooks,visualstudiocode,vim,git,pycharm+all
-# Edit at https://www.toptal.com/developers/gitignore?templates=python,linux,macos,windows,jupyternotebooks,visualstudiocode,vim,git,pycharm+all
+bulk_data
+production_models/**
+!production_models/**/
+!production_models/**/README.md
+logs/
+wandb/
+*.mp4
+*.jpg
+*.pth
+*.png
+temp/
+batch_scripts/
+outputs/
+notebooks/
+
+# I want all JetBrains-specific file to be ignored
+# (gitignore.io defaults allow style and run settings to be included)
+.idea/
+
+# Created by https://www.toptal.com/developers/gitignore/api/git,vim,linux,macos,python,windows,jupyternotebooks,visualstudiocode
+# Edit at https://www.toptal.com/developers/gitignore?templates=git,vim,linux,macos,python,windows,jupyternotebooks,visualstudiocode
 
 ### Git ###
 # Created by git for backups. To disable backups in Git:
@@ -78,94 +97,6 @@ Temporary Items
 # iCloud generated files
 *.icloud
 
-### PyCharm+all ###
-# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider
-# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
-
-# User-specific stuff
-.idea/**/workspace.xml
-.idea/**/tasks.xml
-.idea/**/usage.statistics.xml
-.idea/**/dictionaries
-.idea/**/shelf
-
-# AWS User-specific
-.idea/**/aws.xml
-
-# Generated files
-.idea/**/contentModel.xml
-
-# Sensitive or high-churn files
-.idea/**/dataSources/
-.idea/**/dataSources.ids
-.idea/**/dataSources.local.xml
-.idea/**/sqlDataSources.xml
-.idea/**/dynamic.xml
-.idea/**/uiDesigner.xml
-.idea/**/dbnavigator.xml
-
-# Gradle
-.idea/**/gradle.xml
-.idea/**/libraries
-
-# Gradle and Maven with auto-import
-# When using Gradle or Maven with auto-import, you should exclude module files,
-# since they will be recreated, and may cause churn.  Uncomment if using
-# auto-import.
-# .idea/artifacts
-# .idea/compiler.xml
-# .idea/jarRepositories.xml
-# .idea/modules.xml
-# .idea/*.iml
-# .idea/modules
-# *.iml
-# *.ipr
-
-# CMake
-cmake-build-*/
-
-# Mongo Explorer plugin
-.idea/**/mongoSettings.xml
-
-# File-based project format
-*.iws
-
-# IntelliJ
-out/
-
-# mpeltonen/sbt-idea plugin
-.idea_modules/
-
-# JIRA plugin
-atlassian-ide-plugin.xml
-
-# Cursive Clojure plugin
-.idea/replstate.xml
-
-# SonarLint plugin
-.idea/sonarlint/
-
-# Crashlytics plugin (for Android Studio and IntelliJ)
-com_crashlytics_export_strings.xml
-crashlytics.properties
-crashlytics-build.properties
-fabric.properties
-
-# Editor-based Rest Client
-.idea/httpRequests
-
-# Android studio 3.1+ serialized cache file
-.idea/caches/build_file_checksums.ser
-
-### PyCharm+all Patch ###
-# Ignore everything but code style settings and run configurations
-# that are supposed to be shared within teams.
-
-.idea/*
-
-!.idea/codeStyles
-!.idea/runConfigurations
-
 ### Python ###
 # Byte-compiled / optimized / DLL files
 __pycache__/
@@ -288,7 +219,7 @@ celerybeat.pid
 # Environments
 .env
 .venv
-# env/
+env/
 venv/
 ENV/
 env.bak/
@@ -357,7 +288,7 @@ tags
 
 ### VisualStudioCode ###
 .vscode/*
-# !.vscode/settings.json
+!.vscode/settings.json
 !.vscode/tasks.json
 !.vscode/launch.json
 !.vscode/extensions.json
@@ -400,20 +331,4 @@ $RECYCLE.BIN/
 # Windows shortcuts
 *.lnk
 
-# End of https://www.toptal.com/developers/gitignore/api/python,linux,macos,windows,jupyternotebooks,visualstudiocode,vim,git,pycharm+all
-
-
-bulk_data
-production_models/**
-!production_models/**/
-!production_models/**/README.md
-logs/
-wandb/
-*.mp4
-*.jpg
-*.pth
-*.png
-temp/
-batch_scripts/
-outputs/
-notebooks/
+# End of https://www.toptal.com/developers/gitignore/api/git,vim,linux,macos,python,windows,jupyternotebooks,visualstudiocode

pyproject.toml

@@ -9,8 +9,8 @@ readme = "README.md"
 requires-python = ">=3.13,<3.14"
 dependencies = [
     "numpy>=2.0,<3",
-    "scipy>=1.16.1,<2",
-    "pandas>=2.3,<3",
+    "scipy>=1.16.2,<1.17.0",
+    "pandas>=2.3.3,<2.4.0",
     "matplotlib>=3.10,<4",
     "torch==2.9.*",  # specific version required by torchvision and torchcodec!
     "torchvision==0.24.*",  # specific to torch 2.9
@@ -41,6 +41,9 @@ dependencies = [
     "filelock>=3.20.0,<4",
     "av<16",  # imageio depends on pyav but no pyav 16.0.0 wheel is available for my hardware yet
     "loguru==0.7.*",
+    # PEP561 stub packages for better static analysis
+    "pandas-stubs (>=2.3.3,<2.4.0)",
+    "scipy-stubs (>=1.16.2,<1.17.0)"
     # The following should be installed manually (pip install -e .):
     # flygym & dm_control
     # parallel-video-io

scripts_on_cluster/atomic_batch_extraction/compress_results.run

@@ -0,0 +1,66 @@
+#!/bin/bash -l
+
+#SBATCH --job-name=compress_output
+#SBATCH --nodes 1
+#SBATCH --ntasks 1
+#SBATCH --cpus-per-task 36
+#SBATCH --mem 96GB
+#SBATCH --time 1:00:00
+#SBATCH --partition=standard
+#SBATCH --qos=serial
+#SBATCH --output=logs/compress_output.log
+
+# -------------------------------------------------------------------
+# CRASH CHECK: Ensure SLURM_CPUS_PER_TASK is set
+# -------------------------------------------------------------------
+
+# Get the requested number of CPUs from the Slurm environment
+max_jobs=$SLURM_CPUS_PER_TASK 
+
+# Check if the variable is empty or null.
+if [ -z "$max_jobs" ]; then
+    echo "ERROR: SLURM_CPUS_PER_TASK environment variable is not set." >&2
+    echo "This script must be submitted using 'sbatch' and requires the --cpus-per-task flag in the header." >&2
+    exit 1
+fi
+
+echo "Running up to $max_jobs concurrent tar jobs ($(date))."
+echo "----------------------------------------------"
+
+input_basedir="$HOME/poseforge/bulk_data/pose_estimation/atomic_batches/4variants"
+output_basedir="$HOME/poseforge/bulk_data/pose_estimation/atomic_batches/4variants_tarballs"
+mkdir -p $output_basedir
+
+# Loop through all subdirectories
+job_count=0
+for dir_path in "$input_basedir"/*/; do
+    if [ -d "$dir_path" ]; then
+
+        dir_name=$(basename "$dir_path")
+        output_file="$output_basedir/${dir_name}.tar.gz"
+
+        # Launch the tar command in the background
+        # Note the -C flag to change directory to the parent 
+        # *before* archiving. This keeps the archive clean.
+
+        tar -czf "$output_file" -C "$dir_path/.." "$dir_name" &
+
+        echo "Launched job for $dir_name (PID: $!)"
+        job_count=$((job_count + 1))
+
+        # When we hit the max number of concurrent jobs, we wait for *one*
+        # of the currently running background jobs to complete (wait -n).
+        if [ "$job_count" -ge "$max_jobs" ]; then
+            wait -n 
+            job_count=$((job_count - 1))
+            echo "A background job finished. Resuming launch..."
+        fi
+    fi
+done
+
+# Wait for ALL remaining background jobs to finish
+echo "----------------------------------------------"
+echo "All jobs launched. Waiting for final completion..."
+wait
+
+echo "All parallel tar operations complete ($(date))."