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'bone_{save_heatemapimg_index}.jpg'), heatmap_img[save_heatemapimg_index, 0])
cv2.imwrite(os.path.join(save_path_img, f'pose_{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
Datasets
Models
