Align DRF skeleton preprocessing with upstream heatmap path

datasets/pretreatment_scoliosis_drf.py

@@ -7,7 +7,7 @@ import pickle
 import sys
 from glob import glob
 from pathlib import Path
-from typing import Any, TypedDict, cast
+from typing import Any, Literal, TypedDict, cast
 
 import numpy as np
 import yaml
@@ -34,6 +34,7 @@ JOINT_PAIRS = (
 )
 EPS = 1e-6
 FloatArray = NDArray[np.float32]
+HeatmapReduction = Literal["upstream", "max", "sum"]
 
 
 class SequenceRecord(TypedDict):
@@ -66,6 +67,16 @@ def get_args() -> argparse.Namespace:
         default="configs/drf/pretreatment_heatmap_drf.yaml",
         help="Heatmap preprocessing config used to build the skeleton map branch.",
     )
+    _ = parser.add_argument(
+        "--heatmap_reduction",
+        type=str,
+        choices=["upstream", "max", "sum"],
+        default="upstream",
+        help=(
+            "How to collapse joint/limb heatmaps into one channel. "
+            "'upstream' matches OpenGait at f754f6f..., while 'sum' keeps the paper-literal ablation."
+        ),
+    )
     _ = parser.add_argument(
         "--stats_partition",
         type=str,
@@ -180,7 +191,7 @@ def main() -> None:
         norm_args=heatmap_cfg["norm_args"],
         heatmap_generator_args=heatmap_cfg["heatmap_generator_args"],
         align_args=heatmap_cfg["align_args"],
-        reduction="sum",
+        reduction=cast(HeatmapReduction, args.heatmap_reduction),
     )
 
     pose_paths = iter_pose_paths(args.pose_data_path)