人体姿态估计的热图变成坐标点的两种方案

# import torch.nn as nn
# import torch
#
# import torch.nn.functional as F
#
# class SimpleSelfAttention(nn.Module):
#     def __init__(self, n_in:int, ks=1, sym=False):#, n_out:int):
#         super().__init__()
#         # self.conv = nn.Conv2d(n_in, n_in, kernel_size=1,
#         #                     stride=1, padding=ks//2, bias=False)
#         self.midchannel = n_in//4
#         self.conv = nn.Conv2d(n_in, self.midchannel, ks, padding=ks//2, bias=False)
#         self.conv1 = nn.Conv2d(n_in, 1, ks, padding=ks//2, bias=False)
#         self.convexp = nn.Conv2d(self.midchannel, n_in, ks, padding=ks//2, bias=False)
#         self.gamma = nn.Parameter(torch.tensor([0.]))
#         self.sym = sym
#         self.n_in = n_in
#     def conv1d(self,ni:int, no:int, ks:int=1, stride:int=1, padding:int=0, bias:bool=False):
#         "Create and initialize a `nn.Conv1d` layer with spectral normalization."
#         conv = nn.Conv1d(ni, no, ks, stride=stride, padding=padding, bias=bias)
#         nn.init.kaiming_normal_(conv.weight)
#         if bias: conv.bias.data.zero_()
#         return conv
#     def forward(self,x):
#         if self.sym:
#             # symmetry hack by https://github.com/mgrankin
#             c = self.conv.weight.view(self.n_in,self.n_in)
#             c = (c + c.t())/2
#             self.conv.weight = c.view(self.n_in,self.n_in,1)
#         size = x.size()
#         # print(size)
#         # x = x.view(*size[:2],-1)   # (C,N)
#         x_ = self.conv1(x)
#         x_ = x_.view(size[2],-1)   # (C,N)
#         print('x_=',x_.shape)
#         # changed the order of mutiplication to avoid O(N^2) complexity
#         # (x*xT)*(W*x) instead of (x*(xT*(W*x)))
#         convx = self.conv(x)   # (C,C) * (C,N) = (C,N)   => O(NC^2)
#         print('convx=',convx.shape)
#         # print(x.shape,x_.shape)
#         xxT = torch.mm(x_,x_.permute(1,0).contiguous())   # (C,N) * (N,C) = (C,C)   => O(NC^2)
#         print('xxT=',xxT.shape)
#         o = torch.matmul(xxT, convx)   # (C,C) * (C,N) = (C,N)   => O(NC^2)
#         print('o=',o.shape)
#         o = self.convexp(o)
#         print('o1=', o.shape)
#         o = self.gamma * o + x
#         print('o2=', o.shape)
#         return o.view(*size).contiguous()
#
# b = torch.rand(3,40,32,24)
# aaa = SimpleSelfAttention(40,1)
# aaa(b)


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_modify(heatmaps):
    assert isinstance(heatmaps, torch.Tensor)
    '''
    heatmap: [B, numJoint, H, W]
    '''
    accu_x = heatmaps.sum(dim=2)  # [1, 2, x_dim]
    accu_y = heatmaps.sum(dim=3)

    accu_x = F.softmax(accu_x, 2)
    accu_y = F.softmax(accu_y, 2)

    # accu_x = accu_x * torch.arange(heatmaps.shape[-1]).type(torch.cuda.FloatTensor)
    # accu_y = accu_y * torch.arange(heatmaps.shape[-2]).type(torch.cuda.FloatTensor)

    accu_x = accu_x * torch.arange(heatmaps.shape[-1]).float()
    accu_y = accu_y * torch.arange(heatmaps.shape[-2]).float()

    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__':
    x = torch.from_numpy(np.array([[1,2,3,4],[1,2,3,4]])).type(torch.float32)

    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]
    heatmap = torch.unsqueeze(torch.from_numpy(heatmap), 0)  # [1, numJoint, 64, 64]

    x, y = generate_3d_integral_preds_tensor_modify(heatmaps=heatmap)  #
    print('x=',x)
    print('y=',y)

    print("")


print('***************************************')
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_init(heatmaps, num_joints, x_dim, y_dim):
    assert isinstance(heatmaps, torch.Tensor)
    '''
    heatmap: [B, numJoint, H, W]
    '''
    heatmap = heatmaps.reshape((heatmaps.shape[0], 3, -1))
    heatmap = F.softmax(heatmap, 2)

    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.arange(x_dim).float()
    accu_y = accu_y * torch.arange(y_dim).float()

    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__':
    x = torch.from_numpy(np.array([[1, 2, 3, 4], [1, 2, 3, 4]])).type(torch.float32)

    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) # [1, numJoint, 64, 64]

    x, y = generate_3d_integral_preds_tensor_init(heatmaps=heatmap, num_joints=3, x_dim=32, y_dim=64)
    print('x***********=',x)#
    print('y*************=',y)