SkeletonGait

datasets/pretreatment_heatmap.py

@@ -0,0 +1,712 @@
+import os
+import cv2
+import yaml
+import math
+import torch
+import random
+import pickle
+import argparse
+import numpy as np
+from glob import glob 
+from tqdm import tqdm
+import matplotlib.cm as cm
+import torch.distributed as dist
+from torchvision import transforms as T
+from torch.utils.data import Dataset, DataLoader
+from sklearn.impute import KNNImputer, SimpleImputer
+
+torch.manual_seed(347)
+random.seed(347)
+
+#########################################################################################################
+# The following code is the base class code for generating heatmap.
+#########################################################################################################
+
+class GeneratePoseTarget:
+    """Generate pseudo heatmaps based on joint coordinates and confidence.
+    Required keys are "keypoint", "img_shape", "keypoint_score" (optional),
+    added or modified keys are "imgs".
+    Args:
+        sigma (float): The sigma of the generated gaussian map. Default: 0.6.
+        use_score (bool): Use the confidence score of keypoints as the maximum
+            of the gaussian maps. Default: True.
+        with_kp (bool): Generate pseudo heatmaps for keypoints. Default: True.
+        with_limb (bool): Generate pseudo heatmaps for limbs. At least one of
+            'with_kp' and 'with_limb' should be True. Default: False.
+        skeletons (tuple[tuple]): The definition of human skeletons.
+            Default: ((0, 1), (0, 2), (1, 3), (2, 4), (0, 5), (5, 7), (7, 9),
+                      (0, 6), (6, 8), (8, 10), (5, 11), (11, 13), (13, 15),
+                      (6, 12), (12, 14), (14, 16), (11, 12)),
+            which is the definition of COCO-17p skeletons.
+        double (bool): Output both original heatmaps and flipped heatmaps.
+            Default: False.
+        left_kp (tuple[int]): Indexes of left keypoints, which is used when
+            flipping heatmaps. Default: (1, 3, 5, 7, 9, 11, 13, 15),
+            which is left keypoints in COCO-17p.
+        right_kp (tuple[int]): Indexes of right keypoints, which is used when
+            flipping heatmaps. Default: (2, 4, 6, 8, 10, 12, 14, 16),
+            which is right keypoints in COCO-17p.
+        left_limb (tuple[int]): Indexes of left limbs, which is used when
+            flipping heatmaps. Default: (1, 3, 5, 7, 9, 11, 13, 15),
+            which is left limbs of skeletons we defined for COCO-17p.
+        right_limb (tuple[int]): Indexes of right limbs, which is used when
+            flipping heatmaps. Default: (2, 4, 6, 8, 10, 12, 14, 16),
+            which is right limbs of skeletons we defined for COCO-17p.
+    """
+
+    def __init__(self,
+                 sigma=0.6,
+                 use_score=True,
+                 with_kp=True,
+                 with_limb=False,
+                 skeletons=((0, 1), (0, 2), (1, 3), (2, 4), (0, 5), (5, 7),
+                            (7, 9), (0, 6), (6, 8), (8, 10), (5, 11), (11, 13),
+                            (13, 15), (6, 12), (12, 14), (14, 16), (11, 12)),
+                 double=False,
+                 left_kp=(1, 3, 5, 7, 9, 11, 13, 15),
+                 right_kp=(2, 4, 6, 8, 10, 12, 14, 16),
+                 left_limb=(0, 2, 4, 5, 6, 10, 11, 12),
+                 right_limb=(1, 3, 7, 8, 9, 13, 14, 15),
+                 scaling=1.,
+                 eps= 1e-3,
+                 img_h=64,
+                 img_w = 64):
+
+        self.sigma = sigma
+        self.use_score = use_score
+        self.with_kp = with_kp
+        self.with_limb = with_limb
+        self.double = double
+        self.eps = eps
+
+        assert self.with_kp + self.with_limb == 1, ('One of "with_limb" and "with_kp" should be set as True.')
+        self.left_kp = left_kp
+        self.right_kp = right_kp
+        self.skeletons = skeletons
+        self.left_limb = left_limb
+        self.right_limb = right_limb
+        self.scaling = scaling
+        self.img_h = img_h
+        self.img_w = img_w
+
+    def generate_a_heatmap(self, arr, centers, max_values, point_center):
+        """Generate pseudo heatmap for one keypoint in one frame.
+        Args:
+            arr (np.ndarray): The array to store the generated heatmaps. Shape: img_h * img_w.
+            centers (np.ndarray): The coordinates of corresponding keypoints (of multiple persons). Shape: 1 * 2.
+            max_values (np.ndarray): The max values of each keypoint. Shape: (1, ).
+            point_center: Shape: (1, 2)
+        Returns:
+            np.ndarray: The generated pseudo heatmap.
+        """
+
+        sigma = self.sigma
+        img_h, img_w = arr.shape
+
+        for center, max_value in zip(centers, max_values):
+            if max_value < self.eps:
+                continue
+
+            mu_x, mu_y = center[0], center[1]
+
+            tmp_st_x = int(mu_x - 3 * sigma)
+            tmp_ed_x = int(mu_x + 3 * sigma)
+            tmp_st_y = int(mu_y - 3 * sigma)
+            tmp_ed_y = int(mu_y + 3 * sigma)
+
+            st_x = max(tmp_st_x, 0) 
+            ed_x = min(tmp_ed_x + 1, img_w)
+            st_y = max(tmp_st_y, 0)
+            ed_y = min(tmp_ed_y + 1, img_h)
+            x = np.arange(st_x, ed_x, 1, np.float32)
+            y = np.arange(st_y, ed_y, 1, np.float32)
+
+            # if the keypoint not in the heatmap coordinate system
+            if not (len(x) and len(y)):
+                continue
+            y = y[:, None]
+
+            patch = np.exp(-((x - mu_x)**2 + (y - mu_y)**2) / 2 / sigma**2)
+            patch = patch * max_value
+
+            arr[st_y:ed_y, st_x:ed_x] = np.maximum(arr[st_y:ed_y, st_x:ed_x], patch)
+
+    def generate_a_limb_heatmap(self, arr, starts, ends, start_values, end_values, point_center):
+        """Generate pseudo heatmap for one limb in one frame.
+        Args:
+            arr (np.ndarray): The array to store the generated heatmaps. Shape: img_h * img_w.
+            starts (np.ndarray): The coordinates of one keypoint in the corresponding limbs. Shape: 1 * 2.
+            ends (np.ndarray): The coordinates of the other keypoint in the corresponding limbs. Shape: 1 * 2.
+            start_values (np.ndarray): The max values of one keypoint in the corresponding limbs. Shape: (1, ).
+            end_values (np.ndarray): The max values of the other keypoint in the corresponding limbs. Shape: (1, ).
+        Returns:
+            np.ndarray: The generated pseudo heatmap.
+        """
+
+        sigma = self.sigma
+        img_h, img_w = arr.shape
+
+        for start, end, start_value, end_value in zip(starts, ends, start_values, end_values):
+            value_coeff = min(start_value, end_value)
+            if value_coeff < self.eps:
+                continue
+
+            min_x, max_x = min(start[0], end[0]), max(start[0], end[0])
+            min_y, max_y = min(start[1], end[1]), max(start[1], end[1])
+
+
+            
+            tmp_min_x = int(min_x - 3 * sigma)
+            tmp_max_x = int(max_x + 3 * sigma) 
+            tmp_min_y = int(min_y - 3 * sigma)
+            tmp_max_y = int(max_y + 3 * sigma)
+
+            min_x = max(tmp_min_x, 0)
+            max_x = min(tmp_max_x + 1, img_w)
+            min_y = max(tmp_min_y, 0)
+            max_y = min(tmp_max_y + 1, img_h)
+
+            x = np.arange(min_x, max_x, 1, np.float32)
+            y = np.arange(min_y, max_y, 1, np.float32)
+
+            if not (len(x) and len(y)):
+                continue
+
+            y = y[:, None]
+            x_0 = np.zeros_like(x)
+            y_0 = np.zeros_like(y)
+
+            # distance to start keypoints
+            d2_start = ((x - start[0])**2 + (y - start[1])**2)
+
+            # distance to end keypoints
+            d2_end = ((x - end[0])**2 + (y - end[1])**2)
+
+            # the distance between start and end keypoints.
+            d2_ab = ((start[0] - end[0])**2 + (start[1] - end[1])**2)
+
+            if d2_ab < 1:
+                self.generate_a_heatmap(arr, start[None], start_value[None], point_center)
+                continue
+
+            coeff = (d2_start - d2_end + d2_ab) / 2. / d2_ab
+
+            a_dominate = coeff <= 0
+            b_dominate = coeff >= 1
+            seg_dominate = 1 - a_dominate - b_dominate
+
+            position = np.stack([x + y_0, y + x_0], axis=-1)
+            projection = start + np.stack([coeff, coeff], axis=-1) * (end - start)
+            d2_line = position - projection
+            d2_line = d2_line[:, :, 0]**2 + d2_line[:, :, 1]**2
+            d2_seg = a_dominate * d2_start + b_dominate * d2_end + seg_dominate * d2_line
+
+            patch = np.exp(-d2_seg / 2. / sigma**2)
+            patch = patch * value_coeff
+
+            arr[min_y:max_y, min_x:max_x] = np.maximum(arr[min_y:max_y, min_x:max_x], patch)
+    def generate_heatmap(self, arr, kps, max_values):
+        """Generate pseudo heatmap for all keypoints and limbs in one frame (if
+        needed).
+        Args:
+            arr (np.ndarray): The array to store the generated heatmaps. Shape: V * img_h * img_w.
+            kps (np.ndarray): The coordinates of keypoints in this frame. Shape: 1 * V * 2.
+            max_values (np.ndarray): The confidence score of each keypoint. Shape: 1 * V.
+        Returns:
+            np.ndarray: The generated pseudo heatmap.
+        """
+
+        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)
+
+        if self.with_limb:
+            for i, limb in enumerate(self.skeletons):
+                start_idx, end_idx = limb
+                starts = kps[:, start_idx]
+                ends = kps[:, end_idx]
+
+                start_values = max_values[:, start_idx]
+                end_values = max_values[:, end_idx]
+                self.generate_a_limb_heatmap(arr[i], starts, ends, start_values, end_values, point_center)
+
+    def gen_an_aug(self, pose_data):
+        """Generate pseudo heatmaps for all frames.
+        Args:
+            pose_data (array): [1, T, V, C]
+        Returns:
+            list[np.ndarray]: The generated pseudo heatmaps.
+        """
+
+        all_kps = pose_data[..., :2]
+        kp_shape = pose_data.shape # [1, T, V, 2]
+
+        if pose_data.shape[-1] == 3:
+            all_kpscores = pose_data[..., -1] # [1, T, V]
+        else:
+            all_kpscores = np.ones(kp_shape[:-1], dtype=np.float32)
+
+        
+
+        # scale img_h, img_w and kps
+        img_h = int(self.img_h * self.scaling + 0.5)
+        img_w = int(self.img_w * self.scaling + 0.5)
+        all_kps[..., :2] *= self.scaling
+
+        num_frame = kp_shape[1]
+        num_c = 0
+        if self.with_kp:
+            num_c += all_kps.shape[2]
+        if self.with_limb:
+            num_c += len(self.skeletons)
+        ret = np.zeros([num_frame, num_c, img_h, img_w], dtype=np.float32)
+
+        for i in range(num_frame):
+            # 1, V, C
+            kps = all_kps[:, i]
+            # 1, V
+            kpscores = all_kpscores[:, i] if self.use_score else np.ones_like(all_kpscores[:, i])
+
+            self.generate_heatmap(ret[i], kps, kpscores)
+        return ret
+
+    def __call__(self, pose_data):
+        """
+        pose_data: (T, V, C=3/2)
+        1: means person number
+        """
+        pose_data = pose_data[None,...] # (1, T, V, C=3/2)
+
+        heatmap = self.gen_an_aug(pose_data)
+        
+        if self.double:
+            indices = np.arange(heatmap.shape[1], dtype=np.int64)
+            left, right = (self.left_kp, self.right_kp) if self.with_kp else (self.left_limb, self.right_limb)
+            for l, r in zip(left, right):  # noqa: E741
+                indices[l] = r
+                indices[r] = l
+            heatmap_flip = heatmap[..., ::-1][:, indices]
+            heatmap = np.concatenate([heatmap, heatmap_flip])
+        return heatmap
+
+    def __repr__(self):
+        repr_str = (f'{self.__class__.__name__}('
+                    f'sigma={self.sigma}, '
+                    f'use_score={self.use_score}, '
+                    f'with_kp={self.with_kp}, '
+                    f'with_limb={self.with_limb}, '
+                    f'skeletons={self.skeletons}, '
+                    f'double={self.double}, '
+                    f'left_kp={self.left_kp}, '
+                    f'right_kp={self.right_kp})')
+        return repr_str
+
+class HeatmapToImage:
+    """
+    Convert the heatmap data to image data.
+    """
+    def __init__(self) -> None:
+        self.cmap = cm.gray
+    
+    def __call__(self, heatmaps):
+        """
+        heatmaps: (T, 17, H, W)
+        return images: (T, 1, H, W)
+        """
+        heatmaps = [x.transpose(1, 2, 0) for x in heatmaps]
+        h, w, _ = heatmaps[0].shape
+        newh, neww = int(h), int(w)
+        heatmaps = [np.max(x, axis=-1) for x in heatmaps]
+        heatmaps = [(self.cmap(x)[..., :3] * 255).astype(np.uint8) for x in heatmaps]
+        heatmaps = [cv2.resize(x, (neww, newh)) for x in heatmaps]
+        return np.ascontiguousarray(np.mean(np.array(heatmaps), axis=-1, keepdims=True).transpose(0,3,1,2))
+
+class CenterAndScaleNormalizer:
+
+    def __init__(self, pose_format="coco", use_conf=True, heatmap_image_height=128) -> None:
+        """
+        Parameters:
+        - pose_format (str): Specifies the format of the keypoints. 
+                            This parameter determines how the keypoints are structured and indexed. 
+                            The supported formats are "coco" or "openpose-x" where 'x' can be either 18 or 25, indicating the number of keypoints used by the OpenPose model. 
+        - use_conf (bool): Indicates whether confidence scores.
+        - heatmap_image_height (int): Sets the height (in pixels) for the heatmap images that will be normlization. 
+        """
+        self.pose_format = pose_format
+        self.use_conf = use_conf
+        self.heatmap_image_height = heatmap_image_height
+
+    def __call__(self, data):
+        """
+            Implements step (a) from Figure 2 in the SkeletonGait paper.
+            data: (T, V, C)
+            - T: number of frames
+            - V: number of joints
+            - C: dimensionality, where 2 indicates joint coordinates and 1 indicates the confidence score
+            return data: (T, V, C)
+        """
+
+        if self.use_conf:
+            pose_seq = data[..., :-1]
+            score = np.expand_dims(data[..., -1], axis=-1)
+        else:
+            pose_seq = data[..., :-1]
+        
+        # Hip as the center point
+        if self.pose_format.lower() == "coco":
+            hip  = (pose_seq[:, 11] + pose_seq[:, 12]) / 2. # [t, 2]
+        elif self.pose_format.split('-')[0].lower() == "openpose":
+            hip  = (pose_seq[:, 9] + pose_seq[:, 12]) / 2. # [t, 2]
+        else:
+            raise ValueError(f"Error value for pose_format: {self.pose_format} in CenterAndScale Class.")
+
+        # Center-normalization
+        pose_seq = pose_seq - hip[:, np.newaxis, :]
+
+        # Scale-normalization
+        y_max = np.max(pose_seq[:, :, 1], axis=-1) # [t]
+        y_min = np.min(pose_seq[:, :, 1], axis=-1) # [t]
+        pose_seq *= ((self.heatmap_image_height // 1.5) / (y_max - y_min)[:, np.newaxis, np.newaxis]) # [t, v, 2]
+        
+        pose_seq += self.heatmap_image_height // 2
+        
+        if self.use_conf:
+            pose_seq = np.concatenate([pose_seq, score], axis=-1)
+        return pose_seq
+
+class PadKeypoints:
+    """
+    Pad the keypoints with missing values.
+    """
+
+    def __init__(self, pad_method="knn", use_conf=True) -> None:
+        """
+        pad_method (str): Specifies the method used to pad the missing values.
+                        The supported methods are "knn" and "simple".
+        use_conf (bool): Indicates whether confidence scores.
+        """
+        self.use_conf = use_conf
+        if pad_method.lower() == "knn":
+            self.imputer = KNNImputer(missing_values=0.0, n_neighbors=4, weights="distance", add_indicator=False)
+        elif pad_method.lower()  == "simple":
+            self.imputer = SimpleImputer(missing_values=0.0, strategy='mean',add_indicator=True)
+        else:
+            raise ValueError(f"Error value for padding method: {pad_method}")
+    
+    def __call__(self, raw_data):
+        """
+        raw_data: (T, V, C)
+        - T: number of frames
+        - V: number of joints
+        - C: dimensionality, where 2 indicates joint coordinates and 1 indicates the confidence score
+        return padded_data: (T, V, C)
+        """
+        T, V, C = raw_data.shape
+        if self.use_conf:
+            data = raw_data[..., :-1]
+            score = np.expand_dims(raw_data[..., -1], axis=-1)
+            C = C - 1
+        else:
+            data = raw_data[..., :-1]
+        data = data.reshape((T, V*C))
+        padded_data = self.imputer.fit_transform(data)
+        try:
+            padded_data = padded_data.reshape((T, V, C))
+        except:
+            padded_data = data.reshape((T, V, C))
+        if self.use_conf:
+            padded_data = np.concatenate([padded_data, score], axis=-1)
+        return padded_data
+
+class COCO18toCOCO17:
+    """
+    Transfer COCO18 format (Openpose extracted) to COCO17 format
+    """
+
+    def __init__(self, transfer_to_coco17=True):
+        """
+        transfer_to_coco17 (bool): Indicates whether to transfer the keypoints from COCO18 to COCO17 format.
+        """
+        self.map_dict = {
+                0: 0,# "nose",
+                1: 15,# "left_eye",
+                2: 14,# "right_eye",
+                3: 17,# "left_ear",
+                4: 16,# "right_ear",
+                5: 5,# "left_shoulder",
+                6: 2,# "right_shoulder",
+                7: 6,# "left_elbow",
+                8: 3,# "right_elbow",
+                9: 7,# "left_wrist",
+                10: 4,# "right_wrist",
+                11: 11,# "left_hip",
+                12: 8,# "right_hip",
+                13: 12,# "left_knee",
+                14: 9,# "right_knee",
+                15: 13,# "left_ankle",
+                16: 10,# "right_ankle"
+            }
+        self.transfer = transfer_to_coco17
+    
+    def __call__(self, data):
+
+        """
+        data: (T, 18, C)
+        - T: number of frames
+        - 18: number of joints of COCO18 format
+        - C: dimensionality, where 2 indicates joint coordinates and 1 indicates the confidence score
+        return data: (T, 17, C)
+        """
+
+        if self.transfer:
+            """
+            input data [T, 18, C] coco18 format
+            return data [T, 17, C] coco17 format
+            """
+            T, _, C = data.shape
+            coco17_pkl_data = np.zeros((T, 17, C))
+            for i in range(17):
+                coco17_pkl_data[:,i,:] = data[:,self.map_dict[i],:]
+            return coco17_pkl_data
+        else:
+            return data
+
+class GatherTransform(object):
+    """
+    Gather the different transforms.
+    """
+    def __init__(self, base_transform, transform_bone, transform_joint):
+
+        """
+        base_transform: Some common transform, e.g., COCO18toCOCO17, PadKeypoints, CenterAndScale
+        transform_bone: GeneratePoseTarget for generate bone heatmap
+        transform_joint: GeneratePoseTarget for generate joint heatmap
+        """
+        self.base_transform = base_transform
+        self.transform_bone = transform_bone
+        self.transform_joint = transform_joint
+
+    def __call__(self, pose_data):
+        x = self.base_transform(pose_data)
+        heatmap_bone = self.transform_bone(x) # [T, 1, H, W]
+        heatmap_joint = self.transform_joint(x) # [T, 1, H, W]
+        heatmap = np.concatenate([heatmap_bone, heatmap_joint], axis=1)
+        return heatmap
+
+class HeatmapAlignment():
+    def __init__(self, align=True, final_img_size=64, offset=0, heatmap_image_size=128) -> None:
+        self.align = align
+        self.final_img_size = final_img_size
+        self.offset = offset
+        self.heatmap_image_size = heatmap_image_size
+
+    def center_crop(self, heatmap):
+        """
+        Input: [1, heatmap_image_size, heatmap_image_size]
+        Output: [1, final_img_size, final_img_size]
+        """
+        raw_heatmap = heatmap[0]
+        if self.align: 
+            y_sum = raw_heatmap.sum(axis=1)
+            y_top = (y_sum != 0).argmax(axis=0)
+            y_btm = (y_sum != 0).cumsum(axis=0).argmax(axis=0)
+            height = y_btm - y_top + 1
+            raw_heatmap = raw_heatmap[y_top - self.offset: y_btm + 1 + self.offset, (self.heatmap_image_size // 2) - (height // 2) : (self.heatmap_image_size // 2) + (height // 2) + 1]
+        raw_heatmap = cv2.resize(raw_heatmap, (self.final_img_size, self.final_img_size), interpolation=cv2.INTER_AREA)
+        return raw_heatmap[np.newaxis, :, :] # [1, final_img_size, final_img_size]
+
+    def __call__(self, heatmap_imgs):
+        """
+        heatmap_imgs: (T, 1, raw_size, raw_size)
+        return (T, 1, final_img_size, final_img_size)
+        """
+        heatmap_imgs = heatmap_imgs / 255.
+        heatmap_imgs = np.array([self.center_crop(heatmap_img) for heatmap_img in heatmap_imgs]) 
+        return (heatmap_imgs * 255).astype('uint8')
+
+def GenerateHeatmapTransform(
+    coco18tococo17_args,
+    padkeypoints_args,
+    norm_args,
+    heatmap_generator_args,
+    align_args
+):
+
+    base_transform = T.Compose([
+        COCO18toCOCO17(**coco18tococo17_args),
+        PadKeypoints(**padkeypoints_args), 
+        CenterAndScaleNormalizer(**norm_args), 
+    ])
+
+    heatmap_generator_args["with_limb"] = True
+    heatmap_generator_args["with_kp"] = False
+    transform_bone = T.Compose([
+        GeneratePoseTarget(**heatmap_generator_args), 
+        HeatmapToImage(), 
+        HeatmapAlignment(**align_args) 
+    ])
+
+    heatmap_generator_args["with_limb"] = False
+    heatmap_generator_args["with_kp"] = True
+    transform_joint = T.Compose([
+        GeneratePoseTarget(**heatmap_generator_args), 
+        HeatmapToImage(), 
+        HeatmapAlignment(**align_args) 
+    ])
+
+    transform = T.Compose([
+        GatherTransform(base_transform, transform_bone, transform_joint) # [T, 2, H, W]
+    ])
+
+    return transform
+
+#########################################################################################################
+# The following code is DDP progress codes.
+#########################################################################################################
+class SequentialDistributedSampler(torch.utils.data.sampler.Sampler):
+    """
+    Distributed Sampler that subsamples indicies sequentially,
+    making it easier to collate all results at the end.
+    Even though we only use this sampler for eval and predict (no training),
+    which means that the model params won't have to be synced (i.e. will not hang
+    for synchronization even if varied number of forward passes), we still add extra
+    samples to the sampler to make it evenly divisible (like in `DistributedSampler`)
+    to make it easy to `gather` or `reduce` resulting tensors at the end of the loop.
+    """
+
+    def __init__(self, dataset, batch_size, rank=None, num_replicas=None):
+        if num_replicas is None:
+            if not torch.distributed.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = torch.distributed.get_world_size()
+        if rank is None:
+            if not torch.distributed.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = torch.distributed.get_rank()
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.batch_size = batch_size
+        self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.batch_size / self.num_replicas)) * self.batch_size
+        self.total_size = self.num_samples * self.num_replicas
+
+    def __iter__(self):
+        indices = list(range(len(self.dataset)))
+        # add extra samples to make it evenly divisible
+        indices += [indices[-1]] * (self.total_size - len(indices))
+        # subsample
+        indices = indices[self.rank * self.num_samples : (self.rank + 1) * self.num_samples]
+        return iter(indices)
+
+    def __len__(self):
+        return self.num_samples
+
+
+class TransferDataset(Dataset):
+    def __init__(self, args, generate_heatemap_cfgs) -> None:
+        super().__init__()
+        pose_root = args.pose_data_path
+        sigma = generate_heatemap_cfgs['heatmap_generator_args']['sigma']
+        self.dataset_name = args.dataset_name
+        assert self.dataset_name.lower() in ["sustech1k", "grew", "ccpg", "oumvlp", "ou-mvlp", "gait3d", "casiab", "casiae"], f"Invalid dataset name: {self.dataset_name}"
+        self.save_root = os.path.join(args.save_root, f"{self.dataset_name}_sigma_{sigma}_{args.ext_name}")
+        os.makedirs(self.save_root, exist_ok=True)
+
+        self.heatmap_transform = GenerateHeatmapTransform(**generate_heatemap_cfgs)
+
+        if self.dataset_name.lower() == "sustech1k":
+            self.all_ps_data_paths = sorted(glob(os.path.join(pose_root, "*/*/*/03*.pkl")))
+        else:
+            self.all_ps_data_paths = sorted(glob(os.path.join(pose_root, "*/*/*/*.pkl")))
+
+    def __len__(self):
+        return len(self.all_ps_data_paths)
+    
+    def __getitem__(self, index):
+        pose_path = self.all_ps_data_paths[index]
+        with open(pose_path, "rb") as f:
+            pose_data = pickle.load(f)
+            if self.dataset_name.lower() == "grew":
+                # print(pose_data.shape)
+                pose_data = pose_data[:,2:].reshape(-1, 17, 3)
+        
+        tmp_split = pose_path.split('/')
+
+        heatmap_img = self.heatmap_transform(pose_data) # [T, 2, H, W]
+        
+        save_path_pkl = os.path.join(self.save_root, 'pkl', *tmp_split[-4:-1])
+        os.makedirs(save_path_pkl, exist_ok=True)
+
+        # save some visualization
+        if index < 10:
+            # save images
+            save_path_img = os.path.join(self.save_root, 'images', *tmp_split[-4:-1])
+            os.makedirs(save_path_img, exist_ok=True)
+            # save_heatemapimg_index = random.choice(list(range(heatmap_img.shape[0])))
+            for save_heatemapimg_index in range(heatmap_img.shape[0]):
+                cv2.imwrite(os.path.join(save_path_img, f'pose_{save_heatemapimg_index}.jpg'), heatmap_img[save_heatemapimg_index, 0])
+                cv2.imwrite(os.path.join(save_path_img, f'bone_{save_heatemapimg_index}.jpg'), heatmap_img[save_heatemapimg_index, 1])
+
+        pickle.dump(heatmap_img, open(os.path.join(save_path_pkl, tmp_split[-1]), 'wb'))
+        return None
+
+def mycollate(_):
+    return None
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description='Utility for generating heatmaps from pose data.')
+    parser.add_argument('--pose_data_path', type=str, required=True, help="Path to the root directory containing pose data (.pkl files, ID-level) files.")
+    parser.add_argument('--save_root', type=str, required=True, help="Root directory where generated heatmap .pkl files will be saved (ID-level).")
+    parser.add_argument('--ext_name', type=str, default='', help="Extension name to be appended to the 'save_root' for identification.")
+    parser.add_argument('--dataset_name', type=str, required=True, help="Name of the dataset being preprocessed.")
+    parser.add_argument('--heatemap_cfg_path', type=str, default='configs/skeletongait/pretreatment_heatmap.yaml', help="Path to the heatmap generator configuration file.")
+    parser.add_argument("--local_rank", type=int, default=0, help="Local rank for distributed processing, defaults to 0 for non-distributed setups.")
+    opt = parser.parse_args()
+    return opt
+
+def replace_variables(data, context=None):
+    if context is None:
+        context = {}
+
+    if isinstance(data, dict):
+        for key, value in data.items():
+            data[key] = replace_variables(value, context)
+    elif isinstance(data, list):
+        data = [replace_variables(item, context) for item in data]
+    elif isinstance(data, str):
+        if data.startswith('${') and data.endswith('}'):
+            var_path = data[2:-1].split('.')
+            var_value = context
+            try:
+                for part in var_path:
+                    var_value = var_value[part]
+                return var_value
+            except KeyError:
+                raise ValueError(f"Variable {data} not found in context")
+    return data
+
+if __name__ == "__main__":
+    dist.init_process_group("nccl", init_method='env://')
+    local_rank = torch.distributed.get_rank()
+    world_size = torch.distributed.get_world_size()
+
+    args = get_args()
+
+    # Load the heatmap generator configuration
+    with open(args.heatemap_cfg_path, 'r') as stream:
+        generate_heatemap_cfgs = yaml.safe_load(stream)
+        generate_heatemap_cfgs = replace_variables(generate_heatemap_cfgs, generate_heatemap_cfgs)
+    # Create the dataset
+    dataset = TransferDataset(args, generate_heatemap_cfgs)
+
+    # Create the dataloader
+    dist_sampler = SequentialDistributedSampler(dataset, batch_size=1, rank=local_rank, num_replicas=world_size)
+    dataloader = DataLoader(dataset=dataset, batch_size=1, sampler=dist_sampler, num_workers=8, collate_fn=mycollate)
+    for _, tmp in tqdm(enumerate(dataloader), total=len(dataloader)):
+        pass
+
+