Integral-human-pose-2D

Integral human pose estimation的2D pytorch实现

原始链接：
https://github.com/JimmySuen/integral-human-pose/blob/master/pytorch_projects/common_pytorch/loss/integral.py

import torch.nn.functional as F
import torch
from torch.autograd import Variable
import torch.nn as nn
import numpy as np
import os

def makeGaussian(height, width, sigma = 3, center=None):
        """ Make a square gaussian kernel.
        size is the length of a side of the square
        sigma is full-width-half-maximum, which
        can be thought of as an effective radius.
        """
        x = np.arange(0, width, 1, float)
        y = np.arange(0, height, 1, float)[:, np.newaxis]
        if center is None:
            x0 =  width // 2
            y0 = height // 2
        else:
            x0 = center[0]
            y0 = center[1]
        return 10*np.exp(-4*np.log(2) * ((x-x0)**2 + (y-y0)**2) / sigma**2)
    
def generate_hm(height, width ,joints, maxlenght):
        """ Generate a full Heap Map for every joints in an array
        Args:
            height          : Wanted Height for the Heat Map
            width           : Wanted Width for the Heat Map
            joints          : Array of Joints 15*2
            maxlenght       : Lenght of the Bounding Box
        """
        num_joints = joints.shape[0]
        hm = np.zeros((num_joints, height, width), dtype = np.float32)
        for i in range(num_joints):
            s = int(np.sqrt(maxlenght) * maxlenght * 10 / 4096) + 3
            hm[i,:,:] = makeGaussian(height, width, sigma= s, center= (joints[i,0], joints[i,1]))      
        return hm   


def generate_3d_integral_preds_tensor(heatmaps, num_joints, x_dim, y_dim):
    assert isinstance(heatmaps, torch.Tensor)
	'''
    heatmap: [B, numJoint, H, W]
    '''
    heatmaps = heatmaps.reshape((heatmaps.shape[0], num_joints, y_dim, x_dim))

    accu_x = heatmaps.sum(dim=2)  # [1, 2, x_dim]

    accu_y = heatmaps.sum(dim=3)

    accu_x = accu_x * torch.cuda.comm.broadcast(torch.arange(x_dim).type(torch.cuda.FloatTensor), devices=[accu_x.device.index])[0]
    accu_y = accu_y * torch.cuda.comm.broadcast(torch.arange(y_dim).type(torch.cuda.FloatTensor), devices=[accu_y.device.index])[0]

    accu_x = accu_x.sum(dim=2, keepdim=True)
    accu_y = accu_y.sum(dim=2, keepdim=True)
    return accu_x, accu_y

if __name__ == '__main__':
    GT_xy = np.array([[10,10], [3,8], [16,18]])
    heatmap = generate_hm(64, 32 ,GT_xy, 64)   # [numJoint, 64, 64]

    heatmap = torch.unsqueeze(torch.from_numpy(heatmap), 0).cuda()  # [1, numJoint, 64, 64]

    preds = heatmap.reshape((heatmap.shape[0], 3, -1))
    preds = F.softmax(preds, 2)
    
    x, y = generate_3d_integral_preds_tensor(heatmaps=preds, num_joints=3, x_dim=32, y_dim=64)  # 
    # x,y: [1,numJoint,3]
    preds = torch.cat((x, y), dim=0)  # [2,numJoint,1]
    preds = preds.permute(1, 0, 2)  # [numJoint,2,1]
    preds = preds.squeeze()   ####  [numJoint, 2]
    print('preds ', preds.size(), preds)