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