feat: add scoliosis skeleton experiment tooling

datasets/pretreatment_heatmap.py

@@ -9,6 +9,7 @@ import argparse
 import numpy as np
 from glob import glob 
 from copy import deepcopy
+from collections.abc import Sequence
 from typing import Any, Literal
 from tqdm import tqdm
 import matplotlib.cm as cm
@@ -72,7 +73,8 @@ class GeneratePoseTarget:
                  scaling=1.,
                  eps= 1e-3,
                  img_h=64,
-                 img_w = 64):
+                 img_w = 64,
+                 joint_indices: Sequence[int] | None = None):
 
         self.sigma = sigma
         self.use_score = use_score
@@ -90,6 +92,7 @@ class GeneratePoseTarget:
         self.scaling = scaling
         self.img_h = img_h
         self.img_w = img_w
+        self.joint_indices = tuple(joint_indices) if joint_indices is not None else None
 
     def generate_a_heatmap(self, arr, centers, max_values, point_center):
         """Generate pseudo heatmap for one keypoint in one frame.
@@ -221,9 +224,18 @@ class GeneratePoseTarget:
         point_center = kps.mean(1)
 
         if self.with_kp:
-            num_kp = kps.shape[1]
-            for i in range(num_kp):
-                self.generate_a_heatmap(arr[i], kps[:, i], max_values[:, i], point_center)
+            joint_indices = (
+                tuple(range(kps.shape[1]))
+                if self.joint_indices is None
+                else self.joint_indices
+            )
+            for output_index, joint_index in enumerate(joint_indices):
+                self.generate_a_heatmap(
+                    arr[output_index],
+                    kps[:, joint_index],
+                    max_values[:, joint_index],
+                    point_center,
+                )
 
         if self.with_limb:
             for i, limb in enumerate(self.skeletons):
@@ -261,7 +273,7 @@ class GeneratePoseTarget:
         num_frame = kp_shape[1]
         num_c = 0
         if self.with_kp:
-            num_c += all_kps.shape[2]
+            num_c += all_kps.shape[2] if self.joint_indices is None else len(self.joint_indices)
         if self.with_limb:
             num_c += len(self.skeletons)
         ret = np.zeros([num_frame, num_c, img_h, img_w], dtype=np.float32)