feat: upload project

app.py

@@ -1,7 +1,46 @@
 import gradio as gr
+import time
+import numpy
+import os
+from PIL import Image
+import matplotlib.pyplot as plt
+import torch
+import skimage
+from models.hr_net import hr_w32
+from tool_utils import heatmaps_to_coords,draw_joints
 
-def greet(name):
-    return "Hello " + name + "!!"
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+#Create example list from 'examples/'directory
+example_list=[["./examples/"+example] for example in os.listdir("examples")]
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+def predict(numpy_img):
+    #resize the numpy_image size to (256,256)
+    img_np=skimage.transform.resize(numpy_img,[256,256])
+    #convert numpy_image to tensor
+    img=torch.from_numpy(img_np).permute(2,0,1).unsqueeze(0).float().to(device)
+    #choose model class hr_w32
+    model=hr_w32().to(device)
+    #load weights of model
+    model.load_state_dict(torch.load('./weights/HRNet_epoch20_loss0.000474.pth')['model'])
+    # #set model to pred state
+    model.eval()
+    # #predict the heatmaps of joints
+    start_time=time.time()
+    heatmaps_pred=model(img)
+    heatmaps_pred=heatmaps_pred.double()
+    # #convert output to numpy
+    heatmaps_pred_np=heatmaps_pred.squeeze(0).permute(1,2,0).detach().cpu().numpy()
+    # #heatmaps to joints location
+    coord_joints=heatmaps_to_coords(heatmaps_pred_np,resolu_out=[256,256],prob_threshold=0.1)
+    inference_time=time.time()-start_time
+    inference_time_text="model inference time:{:.4f}s".format(inference_time)
+    # #draw coords on image_np
+    img_rgb=draw_joints(img_np,coord_joints)
+    return img_rgb,inference_time_text
+    
+    
+
+demo=gr.Interface(fn=predict, inputs=gr.Image(),outputs=[gr.Image(type='numpy',width=256,height=256),"text"],examples=example_list)
+
+if __name__=="__main__":
+    demo.launch(show_api=False)

models/hr_net.py

@@ -0,0 +1,163 @@
+import torch
+from torch import nn
+
+from models.modules.blocks.bottleneck import Bottleneck
+from models.modules.stage_module import StageModule
+
+
+def weights_init(m):
+    if isinstance(m, nn.Conv2d):
+        nn.init.normal_(m.weight, std=.01)
+        if m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+    elif isinstance(m, nn.BatchNorm2d):
+        nn.init.constant_(m.weight, 1)
+        if m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+            
+
+class HRNet(nn.Module):
+
+    def __init__(self, c=48, nof_joints=16, bn_momentum=.1):
+        super(HRNet, self).__init__()
+
+        # (b,3,y,x) -> (b,64,y,x)
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3,
+                               stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64, momentum=bn_momentum)
+        self.conv2 = nn.Conv2d(64, 64, kernel_size=3,
+                               stride=2, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(64, momentum=bn_momentum)
+        self.relu = nn.ReLU(inplace=True)
+
+        # (b,64,y,x) -> (b,256,y,x)
+        downsample = nn.Sequential(
+            nn.Conv2d(64, 256, kernel_size=1, stride=1, bias=False),
+            nn.BatchNorm2d(256),
+        )
+        self.layer1 = nn.Sequential(
+            Bottleneck(64, 64, downsample=downsample),
+            Bottleneck(256, 64),
+            Bottleneck(256, 64),
+            Bottleneck(256, 64),
+        )
+
+        # (b,256,y,x) ---+---> (b,c,y,x)
+        #                +---> (b,c*2,y/2,x/2)
+        self.transition1 = nn.ModuleList([
+            nn.Sequential(
+                nn.Conv2d(256, c, kernel_size=3,
+                          stride=1, padding=1, bias=False),
+                nn.BatchNorm2d(c),
+                nn.ReLU(inplace=True),
+            ),
+            nn.Sequential(nn.Sequential(
+                nn.Conv2d(256, c * 2, kernel_size=3,
+                          stride=2, padding=1, bias=False),
+                nn.BatchNorm2d(c * 2),
+                nn.ReLU(inplace=True),
+            ))
+        ])
+
+        # StageModule中每个分枝发生了融合
+        # (b,c,y,x) ------+---> (b,c,y,x)
+        # (b,c*2,y/2,x/2) +---> (b,c*2,y/2,x/2)
+        self.stage2 = nn.Sequential(
+            StageModule(stage=2, output_branches=2, c=c, bn_momentum=bn_momentum)
+        )
+
+        # (b,c,y,x) ----------> (b,c,y,x)
+        # (b,c*2,y/2,x/2) +---> (b,c*2,y/2,x/2)
+        #                 +---> (b,c*4,y/4,x/4)
+        self.transition2 = nn.ModuleList([
+            nn.Sequential(),
+            nn.Sequential(),
+            nn.Sequential(nn.Sequential(
+                nn.Conv2d(c * 2, c * 4, kernel_size=3,
+                          stride=2, padding=1, bias=False),
+                nn.BatchNorm2d(c * 4),
+                nn.ReLU(inplace=True),
+            ))
+        ])
+
+        # (b,c,y,x) ------++++---> (b,c,y,x)
+        # (b,c*2,y/2,x/2) ++++---> (b,c*2,y/2,x/2)
+        # (b,c*4,y/4,x/4) ++++---> (b,c*4,y/4,x/4)
+        self.stage3 = nn.Sequential(
+            StageModule(stage=3, output_branches=3, c=c, bn_momentum=bn_momentum),
+            StageModule(stage=3, output_branches=3, c=c, bn_momentum=bn_momentum),
+            StageModule(stage=3, output_branches=3, c=c, bn_momentum=bn_momentum),
+            StageModule(stage=3, output_branches=3, c=c, bn_momentum=bn_momentum),
+        )
+
+        # (b,c,y,x) ----------> (b,c,y,x)
+        # (b,c*2,y/2,x/2) ----> (b,c*2,y/2,x/2)
+        # (b,c*4,y/4,x/4) +---> (b,c*4,y/4,x/4)
+        #                 +---> (b,c*8,y/8,x/8)
+        self.transition3 = nn.ModuleList([
+            nn.Sequential(),  # None,   - Used in place of "None" because it is callable
+            nn.Sequential(),  # None,   - Used in place of "None" because it is callable
+            nn.Sequential(),  # None,   - Used in place of "None" because it is callable
+            nn.Sequential(nn.Sequential(  # Double Sequential to fit with official pretrained weights
+                nn.Conv2d(c * 4, c * 8, kernel_size=3,
+                          stride=2, padding=1, bias=False),
+                nn.BatchNorm2d(c * 8),
+                nn.ReLU(inplace=True),
+            )),
+        ])
+
+        # (b,c,y,x) ------+++---> (b,c,y,x)
+        # (b,c*2,y/2,x/2) +++---> (b,c*2,y/2,x/2)
+        # (b,c*4,y/4,x/4) +++---> (b,c*4,y/4,x/4)
+        # (b,c*8,y/8,x/8) +++---> (b,c*8,y/8,x/8)
+        self.stage4 = nn.Sequential(
+            StageModule(stage=4, output_branches=4, c=c, bn_momentum=bn_momentum),
+            StageModule(stage=4, output_branches=4, c=c, bn_momentum=bn_momentum),
+            StageModule(stage=4, output_branches=1, c=c, bn_momentum=bn_momentum),
+        )
+
+        # 取最高分辨率的结果
+        # (b,c,y,x) -> (b,nof_joints*2,y,x)
+        self.final_layer = nn.Conv2d(c, nof_joints, kernel_size=1, stride=1)
+
+        self.apply(weights_init)
+
+    def forward(self, x):
+        x = self.relu(self.bn1(self.conv1(x)))
+        x = self.relu(self.bn2(self.conv2(x)))
+
+        x = self.layer1(x)
+        x = [trans(x) for trans in self.transition1]
+
+        x = self.stage2(x)
+        x = [
+            self.transition2[0](x[0]),
+            self.transition2[1](x[1]),
+            self.transition2[2](x[1]),
+        ]
+
+        x = self.stage3(x)
+        x = [
+            self.transition3[0](x[0]),
+            self.transition3[1](x[1]),
+            self.transition3[2](x[2]),
+            self.transition3[3](x[2]),
+        ]
+
+        x = self.stage4(x)
+
+        x = x[0]
+        out = self.final_layer(x)
+
+        return out
+
+def hr_w32():
+    return HRNet(32)
+
+if __name__ == '__main__':
+    import torch
+
+    model = hr_w32()
+    x = torch.randn(1,3,256,256)
+    output = model(x)
+    print(output.size())

models/modules/blocks/basic_block.py

@@ -0,0 +1,37 @@
+from torch import nn
+
+
+class BasicBlock(nn.Module):
+    """
+    (b,c,y,x) -> (b,c,y,x)
+    """
+    expansion = 1
+
+    def __init__(self, planes, bn_momentum=.1):
+        super(BasicBlock, self).__init__()
+
+        self.conv1 = nn.Conv2d(planes, planes, kernel_size=3,
+                               stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes, momentum=bn_momentum)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
+                               stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes, momentum=bn_momentum)
+
+    def forward(self, x):
+        residual = x
+
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+
+        out += residual
+        return self.relu(out)
+
+
+if __name__ == '__main__':
+    import torch
+
+    model = BasicBlock(256)
+    x = torch.randn(1, 256, 128, 128)
+    print(model(x).size())  # torch.Size([1,256,128,128])

models/modules/blocks/bottleneck.py

@@ -0,0 +1,57 @@
+from torch import nn
+
+
+class Bottleneck(nn.Module):
+    """
+    (b,c_in,y,x) -> (b,4*c_out,y,x)
+    """
+
+    expansion = 4
+
+    def __init__(self, inplanes, planes, downsample=None, bn_momentum=.1):
+        super(Bottleneck, self).__init__()
+
+        self.conv1 = nn.Conv2d(inplanes, planes,
+                               kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes, momentum=bn_momentum)
+
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
+                               stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes, momentum=bn_momentum)
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion,
+                               kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion,
+                                  momentum=bn_momentum)
+
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+
+    def forward(self, x):
+        residual = x
+
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.relu(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        return self.relu(out)
+
+
+if __name__ == '__main__':
+    import torch
+
+    downsample = nn.Sequential(
+        nn.Conv2d(64, 256, kernel_size=1, stride=1, bias=False),
+        nn.BatchNorm2d(256),
+    )
+    model = Bottleneck(64, 64, downsample=downsample)
+    x = torch.randn(1, 64, 128, 128)
+    print(model(x).size()) # torch.Size([1,256,128,128])
+
+    model = Bottleneck(256,64)
+    x = torch.randn(1,256,128,128)
+    print(model(x).size()) # torch.Size([2,256,128,128])

models/modules/stage_module.py

@@ -0,0 +1,104 @@
+from torch import nn
+
+from models.modules.blocks.basic_block import BasicBlock
+
+
+class StageModule(nn.Module):
+
+    def __init__(self, stage, output_branches, c, bn_momentum):
+        super(StageModule, self).__init__()
+
+        self.stage = stage
+        self.output_branches = output_branches
+
+        # 得到stage对应数量的分枝
+        # 例如stage=3，c=32时
+        # i = 0,1,2
+        # i = 0 -> 4*BasicBlock(32)
+        # i = 1 -> 4*BasicBlock(64)
+        # i = 2 -> 4*BasicBlock(128)
+        #
+        # -+--- 4*BasicBlock(32) ---->
+        #  +--- 4*BasicBlock(64) ---->
+        #  +--- 4*BasicBlock(128) --->
+        self.branches = nn.ModuleList()
+        for i in range(self.stage):
+            w = c * (2**i)
+            branch = nn.Sequential(
+                BasicBlock(w, bn_momentum=bn_momentum),
+                BasicBlock(w, bn_momentum=bn_momentum),
+                BasicBlock(w, bn_momentum=bn_momentum),
+                BasicBlock(w, bn_momentum=bn_momentum),
+            )
+            self.branches.append(branch)
+
+        self.fuse_layers = nn.ModuleList()
+
+        # 得到i*j个输出分枝，其中第(i,j)个输出分枝代表第j个分枝向第i个输出变换的输出分枝
+        # i<j，则输出分枝的通道数小于分枝i的通道数，作上采样
+        # i>j，则输出分枝的通道数大于分枝i的通道数，作下采样
+        #                     +---output branch 0(c=32)---->
+        #                     +(upsample)
+        # ---branch 1(c=64)---+---output branch 1(c=64)---->
+        #                     +(downsample)
+        #                     +---output branch 2(c=128)--->
+        # 对于每一个输出分枝i
+        for i in range(self.output_branches):
+            self.fuse_layers.append(nn.ModuleList())
+
+            # 对于每一个分枝j
+            for j in range(self.stage):
+
+                # 如果分枝与输出分枝相对应，直接输出
+                if i == j:
+                    self.fuse_layers[-1].append(nn.Sequential())
+
+                # 如果输出分枝编号小于分枝编号，则上采样后输出
+                elif i < j:
+                    self.fuse_layers[-1].append(nn.Sequential(
+                        nn.Conv2d(c * (2**j), c * (2**i), kernel_size=1,
+                                  stride=1, bias=False),
+                        nn.BatchNorm2d(c * (2**i)),
+                        nn.Upsample(scale_factor=(2.**(j-i))),
+                    ))
+
+                # 如果输出分枝编号大于分枝编号，则下采样后输出
+                elif i > j:
+                    ops = []
+                    for _ in range(i - j - 1):
+                        ops.append(nn.Sequential(
+                            nn.Conv2d(c * (2**j), c * (2**j), kernel_size=3,
+                                      stride=2, padding=1, bias=False),
+                            nn.BatchNorm2d(c * (2**j)),
+                            nn.ReLU(inplace=True),
+                        ))
+                    ops.append(nn.Sequential(
+                        nn.Conv2d(c * (2**j), c * (2**i), kernel_size=3,
+                                  stride=2, padding=1, bias=False),
+                        nn.BatchNorm2d(c * (2**i)),
+                    ))
+                    self.fuse_layers[-1].append(nn.Sequential(*ops))
+
+            self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        # 将x经过每个分枝
+        x = [branch(b) for branch, b in zip(self.branches, x)]
+
+        x_fused = []
+        # 对于每个输出分枝
+        for i in range(len(self.fuse_layers)):
+            # 对于每个分枝
+            for j in range(len(self.branches)):
+                # 如果是第0个分枝，则将经过第0个分枝的x经过第i个输出分枝
+                if j == 0:
+                    x_fused.append(self.fuse_layers[i][0](x[0]))
+                # 否则，将经过第j个分枝的x经过第i个输出分枝，与之前第i个输出分枝的结果相加
+                else:
+                    x_fused[i] = x_fused[i] + self.fuse_layers[i][j](x[j])
+
+        # 每个输出分枝的结果经过ReLU
+        for i in range(len(x_fused)):
+            x_fused[i] = self.relu(x_fused[i])
+
+        return x_fused

models/modules/stem.py

@@ -0,0 +1,29 @@
+from torch import nn
+
+
+class Stem(nn.Module):
+    """
+    Stem模块进行1/4的下采样，并将通道数变为64
+    (b,3,y,x) -> (b,64,y/4,x/4)
+    """
+    def __init__(self, bn_momentum=.1):
+        super(Stem, self).__init__()
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3,
+                               stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64, momentum=bn_momentum)
+        self.conv2 = nn.Conv2d(64, 64, kernel_size=3,
+                               stride=2, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(64, momentum=bn_momentum)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        out = self.bn1(self.conv1(x))
+        out = self.bn2(self.conv2(out))
+        return self.relu(out)
+
+if __name__ == '__main__':
+    import torch
+
+    model = Stem()
+    x = torch.randn(1,3,128,64)
+    print(model(x).size()) # torch.Size([1,64,32,16])

requirements.txt

@@ -0,0 +1,17 @@
+Cython                    3.0.5
+gradio                    4.8.0
+gradio_client             0.7.1
+huggingface-hub           0.19.4
+imageio                   2.33.0
+numpy                     1.24.3
+opencv-python             4.8.1.78
+opendatalab               0.0.10
+Pillow                    10.0.1
+pip                       23.3
+pycocotools               2.0.7
+scikit-image              0.21.0
+scipy                     1.10.1
+torch                     1.12.0+cu113
+torchaudio                0.12.0+cu113
+torchvision               0.13.0+cu113
+tqdm                      4.65.2

tool_utils.py

@@ -0,0 +1,372 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.image as mpimg
+import cv2
+import skimage
+import torch
+from PIL import Image
+
+joints = [
+    'left ankle',
+    'left knee',
+    'left hip',
+    'right hip',
+    'right knee',
+    'right ankle',
+    'belly',
+    'chest',
+    'neck',
+    'head',
+    'left wrist',
+    'left elbow',
+    'left shoulder',
+    'right shoulder',
+    'right elbow',
+    'right wrist'
+]
+
+
+def generate_heatmap(heatmap, pt, sigma=(33, 33), sigma_valu=7):
+    '''
+    :param heatmap: should be a np zeros array with shape (H,W) (only i channel), not (H,W,1)
+    :param pt: point coords, np array
+    :param sigma: should be a tuple with odd values (obsolete)
+    :param sigma_valu: vaalue for gaussian blur
+    :return: a np array of one joint heatmap with shape (H,W)
+
+    This function is obsolete, use 'generate_heatmaps()' instead.
+    '''
+    heatmap[int(pt[1])][int(pt[0])] = 1
+    # heatmap = cv2.GaussianBlur(heatmap, sigma, 0)  #(H,W,1) -> (H,W)
+    heatmap = skimage.filters.gaussian(
+        heatmap, sigma=sigma_valu)  # (H,W,1) -> (H,W)
+    am = np.amax(heatmap)
+    heatmap = heatmap/am
+    return heatmap
+
+
+def generate_heatmaps(img, pts, sigma=(33, 33), sigma_valu=7):
+    '''
+    :param img: np arrray img, (H,W,C)
+    :param pts: joint points coords, np array, same resolu as img
+    :param sigma: should be a tuple with odd values (obsolete)
+    :param sigma_valu: vaalue for gaussian blur
+    :return: np array heatmaps, (H,W,num_pts)
+    '''
+    H, W = img.shape[0], img.shape[1]
+    num_pts = pts.shape[0]
+    heatmaps = np.zeros((H, W, num_pts))
+    for i, pt in enumerate(pts):
+        # Filter unavailable heatmaps
+        if pt[0] == 0 and pt[1] == 0:
+            continue
+        # Filter some points out of the image
+        if pt[0] >= W:
+            pt[0] = W-1
+        if pt[1] >= H:
+            pt[1] = H-1
+        heatmap = heatmaps[:, :, i]
+        heatmap[int(pt[1])][int(pt[0])] = 1
+        # heatmap = cv2.GaussianBlur(heatmap, sigma, 0)  #(H,W,1) -> (H,W)
+        heatmap = skimage.filters.gaussian(
+            heatmap, sigma=sigma_valu)  # (H,W,1) -> (H,W)
+        am = np.amax(heatmap)
+        heatmap = heatmap / am
+        heatmaps[:, :, i] = heatmap
+    return heatmaps
+
+
+def load_image(path_image):
+    img = mpimg.imread(path_image)
+    # Return a np array (H,W,C)
+    return img
+
+
+def crop(img, ele_anno, use_randscale=True, use_randflipLR=False, use_randcolor=False):
+    '''
+    :param img: np array of the origin image, (H,W,C)
+    :param ele_anno: one element of json annotation
+    :return: img_crop, ary_pts_crop, c_crop after cropping
+    '''
+
+    H, W = img.shape[0], img.shape[1]
+    s = ele_anno['scale_provided']
+    c = ele_anno['objpos']
+
+    # Adjust center and scale
+    if c[0] != -1:
+        c[1] = c[1] + 15 * s
+        s = s * 1.25
+    ary_pts = np.array(ele_anno['joint_self'])  # (16, 3)
+    ary_pts_temp = ary_pts[np.any(ary_pts != [0, 0, 0], axis=1)]
+
+    if use_randscale:
+        scale_rand = np.random.uniform(low=1.0, high=3.0)
+    else:
+        scale_rand = 1
+
+    W_min = max(np.amin(ary_pts_temp, axis=0)[0] - s * 15 * scale_rand, 0)
+    H_min = max(np.amin(ary_pts_temp, axis=0)[1] - s * 15 * scale_rand, 0)
+    W_max = min(np.amax(ary_pts_temp, axis=0)[0] + s * 15 * scale_rand, W)
+    H_max = min(np.amax(ary_pts_temp, axis=0)[1] + s * 15 * scale_rand, H)
+    W_len = W_max - W_min
+    H_len = H_max - H_min
+    window_len = max(H_len, W_len)
+    pad_updown = (window_len - H_len)/2
+    pad_leftright = (window_len - W_len)/2
+
+    # Calculate 4 corner position
+    W_low = max((W_min - pad_leftright), 0)
+    W_high = min((W_max + pad_leftright), W)
+    H_low = max((H_min - pad_updown), 0)
+    H_high = min((H_max + pad_updown), H)
+
+    # Update joint points and center
+    ary_pts_crop = np.where(
+        ary_pts == [0, 0, 0], ary_pts, ary_pts - np.array([W_low, H_low, 0]))
+    c_crop = c - np.array([W_low, H_low])
+
+    img_crop = img[int(H_low):int(H_high), int(W_low):int(W_high), :]
+
+    # Pad when H, W different
+    H_new, W_new = img_crop.shape[0], img_crop.shape[1]
+    window_len_new = max(H_new, W_new)
+    pad_updown_new = int((window_len_new - H_new)/2)
+    pad_leftright_new = int((window_len_new - W_new)/2)
+
+    # ReUpdate joint points and center (because of the padding)
+    ary_pts_crop = np.where(ary_pts_crop == [
+                            0, 0, 0], ary_pts_crop, ary_pts_crop + np.array([pad_leftright_new, pad_updown_new, 0]))
+    c_crop = c_crop + np.array([pad_leftright_new, pad_updown_new])
+
+    img_crop = cv2.copyMakeBorder(img_crop, pad_updown_new, pad_updown_new,
+                                  pad_leftright_new, pad_leftright_new, cv2.BORDER_CONSTANT, value=0)
+
+    # change dtype and num scale
+    img_crop = img_crop / 255.
+    img_crop = img_crop.astype(np.float64)
+
+    if use_randflipLR:
+        flip = np.random.random() > 0.5
+        # print('rand_flipLR', flip)
+        if flip:
+            # (H,W,C)
+            img_crop = np.flip(img_crop, 1)
+            # Calculate flip pts, remember to filter [0,0] which is no available heatmap
+            ary_pts_crop = np.where(ary_pts_crop == [0, 0, 0], ary_pts_crop,
+                                    [window_len_new, 0, 0] + ary_pts_crop * [-1, 1, 0])
+            c_crop = [window_len_new, 0] + c_crop * [-1, 1]
+            # Rearrange pts
+            ary_pts_crop = np.concatenate(
+                (ary_pts_crop[5::-1], ary_pts_crop[6:10], ary_pts_crop[15:9:-1]))
+
+    if use_randcolor:
+        randcolor = np.random.random() > 0.5
+        # print('rand_color', randcolor)
+        if randcolor:
+            img_crop[...,
+                     0] *= np.clip(np.random.uniform(low=0.8, high=1.2), 0., 1.)
+            img_crop[...,
+                     1] *= np.clip(np.random.uniform(low=0.8, high=1.2), 0., 1.)
+            img_crop[...,
+                     2] *= np.clip(np.random.uniform(low=0.8, high=1.2), 0., 1.)
+
+    return img_crop, ary_pts_crop, c_crop
+
+
+def change_resolu(img, pts, c, resolu_out=(256, 256)):
+    '''
+    :param img: np array of the origin image
+    :param pts: joint points np array corresponding to the image, same resolu as img
+    :param c: center
+    :param resolu_out: a list or tuple
+    :return: img_out, pts_out, c_out under resolu_out
+    '''
+    H_in = img.shape[0]
+    W_in = img.shape[1]
+    H_out = resolu_out[0]
+    W_out = resolu_out[1]
+    H_scale = H_in/H_out
+    W_scale = W_in/W_out
+
+    pts_out = pts/np.array([W_scale, H_scale, 1])
+    c_out = c/np.array([W_scale, H_scale])
+    img_out = skimage.transform.resize(img, tuple(resolu_out))
+
+    return img_out, pts_out, c_out
+
+
+def heatmaps_to_coords(heatmaps, resolu_out=[64, 64], prob_threshold=0.2):
+    '''
+    :param heatmaps: tensor with shape (64,64,16)
+    :param resolu_out: output resolution list
+    :return coord_joints: np array, shape (16,2)
+    '''
+
+    num_joints = heatmaps.shape[2]
+    # Resize
+    heatmaps = skimage.transform.resize(heatmaps, tuple(resolu_out))
+
+    coord_joints = np.zeros((num_joints, 3))
+    for i in range(num_joints):
+        heatmap = heatmaps[..., i]
+        max = np.max(heatmap)
+        # Only keep points larger than a threshold
+        if max >= prob_threshold:
+            idx = np.where(heatmap == max)
+            H = idx[0][0]
+            W = idx[1][0]
+        else:
+            H = 0
+            W = 0
+        coord_joints[i] = [W, H, max]
+    return coord_joints
+
+
+def show_heatmaps(img, heatmaps, c=np.zeros((2)), num_fig=1):
+    '''
+    :param img: np array (H,W,3)
+    :param heatmaps: np array (H,W,num_pts)
+    :param c: center, np array (2,)
+    '''
+    H, W = img.shape[0], img.shape[1]
+
+    if heatmaps.shape[0] != H:
+        heatmaps = skimage.transform.resize(heatmaps, (H, W))
+
+    plt.figure(num_fig)
+    for i in range(heatmaps.shape[2] + 1):
+        plt.subplot(4, 5, i + 1)
+        if i == 0:
+            plt.title('Origin')
+        else:
+            plt.title(joints[i-1])
+
+        if i == 0:
+            plt.imshow(img)
+        else:
+            plt.imshow(heatmaps[:, :, i - 1])
+
+        plt.axis('off')
+    plt.subplot(4, 5, 20)
+    plt.axis('off')
+    plt.show()
+
+
+def heatmap2rgb(heatmap):
+    """
+    : heatmap: (h,w)
+    """
+
+    heatmap = heatmap.detach().cpu().numpy()
+
+    # plt.figure(figsize=(1,1))
+    # plt.axis('off')
+    # plt.imshow(heatmap)
+    # plt.savefig('tmp/tmp.jpg', bbox_inches='tight', pad_inches=0, dpi=70)
+    # plt.close()
+    # plt.clf()
+
+    # img = Image.open('tmp/tmp.jpg')
+    cm = plt.get_cmap('jet')
+    normed_data = (heatmap - np.min(heatmap)) / (np.max(heatmap) - np.min(heatmap + 1e-8))
+    mapped_data = cm(normed_data)
+
+    # (h,w,c)
+    # img = np.array(img)
+    img = np.array(mapped_data)
+    img = img[:,:,:3]
+    img = torch.tensor(img).permute(2, 0, 1)
+    
+    return img
+
+
+def heatmaps2rgb(heatmaps):
+    """
+    : heatmaps: (b,h,w)
+    """
+
+    out_imgs = []
+    for heatmap in heatmaps:
+        out_imgs.append(heatmap2rgb(heatmap))
+
+    return torch.stack(out_imgs)
+
+
+# def draw_joints(img, pts):
+#     scores = pts[:,2]
+#     pts = np.array(pts).astype(int)
+
+#     for i in range(pts.shape[0]):
+#         if pts[i, 0] != 0 and pts[i, 1] != 0:
+#             img = cv2.circle(img, (pts[i, 0], pts[i, 1]), radius=3,
+#                              color=(255, 0, 0), thickness=-1)
+#             print('img',img.max(),img.min())
+#             # img = cv2.putText(img, f'{joints[i]}: {scores[i]:.2f}', (
+#             #     pts[i, 0]+5, pts[i, 1]-5), cv2.FONT_HERSHEY_SIMPLEX, .25, (255, 0, 0))
+
+#     # Left arm
+#     for i in range(10, 13-1):
+#         if pts[i, 0] != 0 and pts[i, 1] != 0 and pts[i+1, 0] != 0 and pts[i+1, 1] != 0:
+#             img = cv2.line(img, (pts[i, 0], pts[i, 1]), (pts[i+1, 0],
+#                            pts[i+1, 1]), color=(255, 0, 0), thickness=1)
+
+#     # Right arm
+#     for i in range(13, 16-1):
+#         if pts[i, 0] != 0 and pts[i, 1] != 0 and pts[i+1, 0] != 0 and pts[i+1, 1] != 0:
+#             img = cv2.line(img, (pts[i, 0], pts[i, 1]), (pts[i+1, 0],
+#                            pts[i+1, 1]), color=(255, 0, 0), thickness=1)
+
+#     # Left leg
+#     for i in range(0, 3-1):
+#         if pts[i, 0] != 0 and pts[i, 1] != 0 and pts[i+1, 0] != 0 and pts[i+1, 1] != 0:
+#             img = cv2.line(img, (pts[i, 0], pts[i, 1]), (pts[i+1, 0],
+#                            pts[i+1, 1]), color=(255, 0, 0), thickness=1)
+#     # Right leg
+#     for i in range(3, 6-1):
+#         if pts[i, 0] != 0 and pts[i, 1] != 0 and pts[i+1, 0] != 0 and pts[i+1, 1] != 0:
+#             img = cv2.line(img, (pts[i, 0], pts[i, 1]), (pts[i+1, 0],
+#                            pts[i+1, 1]), color=(255, 0, 0), thickness=1)
+
+#     # Body
+#     for i in range(6, 10-1):
+#         if pts[i, 0] != 0 and pts[i, 1] != 0 and pts[i+1, 0] != 0 and pts[i+1, 1] != 0:
+#             img = cv2.line(img, (pts[i, 0], pts[i, 1]), (pts[i+1, 0],
+#                            pts[i+1, 1]), color=(255, 0, 0), thickness=1)
+
+#     if pts[2, 0] != 0 and pts[2, 1] != 0 and pts[3, 0] != 0 and pts[3, 1] != 0:
+#         img = cv2.line(img, (pts[2, 0], pts[2, 1]), (pts[2+1, 0],
+#                        pts[2+1, 1]), color=(255, 0, 0), thickness=1)
+#     if pts[12, 0] != 0 and pts[12, 1] != 0 and pts[13, 0] != 0 and pts[13, 1] != 0:
+#         img = cv2.line(img, (pts[12, 0], pts[12, 1]), (pts[12+1, 0],
+#                        pts[12+1, 1]), color=(255, 0, 0), thickness=1)
+
+#     return img
+def draw_joints(img, pts):
+    # Convert the image to the range [0, 255] for visualization
+    img_visualization = (img * 255).astype(np.uint8)
+
+    # Draw lines for the body parts
+    for i in range(10, 13 - 1):
+        draw_line(img_visualization, pts[i], pts[i + 1])
+
+    for i in range(13, 16 - 1):
+        draw_line(img_visualization, pts[i], pts[i + 1])
+
+    for i in range(0, 3 - 1):
+        draw_line(img_visualization, pts[i], pts[i + 1])
+
+    for i in range(3, 6 - 1):
+        draw_line(img_visualization, pts[i], pts[i + 1])
+
+    for i in range(6, 10 - 1):
+        draw_line(img_visualization, pts[i], pts[i + 1])
+
+    draw_line(img_visualization, pts[2], pts[3])
+    draw_line(img_visualization, pts[12], pts[13])
+
+    return img_visualization / 255.0
+
+def draw_line(img, pt1, pt2):
+    if pt1[0] != 0 and pt1[1] != 0 and pt2[0] != 0 and pt2[1] != 0:
+        cv2.line(img, (int(pt1[0]), int(pt1[1])), (int(pt2[0]), int(pt2[1])), color=(255, 0, 0), thickness=1)