【AI达人创造营第二期】TrtPose复现、手动转TRT并训练动物姿态

class ResNetBackbone(nn.Layer):
    def __init__(self, resnet):
        super(ResNetBackbone, self).__init__()
        self.resnet = resnet 
    def forward(self, x):
        x = self.resnet.conv1(x)
        x = self.resnet.bn1(x)
        x = self.resnet.relu(x)
        x = self.resnet.maxpool(x)
        x = self.resnet.layer1(x) # /4
        x = self.resnet.layer2(x) # /8
        x = self.resnet.layer3(x) # /16
        x = self.resnet.layer4(x) # /32
        return xclass UpsampleCBR(nn.Sequential):
    def __init__(self, input_channels, output_channels, count=1, num_flat=0):
        layers = []        for i in range(count):            if i == 0:
                inch = input_channels            else:
                inch = output_channels
                
            layers += [
                nn.Conv2DTranspose(inch, output_channels, kernel_size=4, stride=2, padding=1),
                nn.BatchNorm2D(output_channels),
                nn.ReLU()
            ]            for i in range(num_flat):
                layers += [
                    nn.Conv2D(output_channels, output_channels, kernel_size=3, stride=1, padding=1),
                    nn.BatchNorm2D(output_channels),
                    nn.ReLU()
                ]            
        super(UpsampleCBR, self).__init__(*layers)class CmapPafHeadAttention(nn.Layer):
    def __init__(self, input_channels, cmap_channels, paf_channels, upsample_channels=256, num_upsample=0, num_flat=0):
        super(CmapPafHeadAttention, self).__init__()
        self.cmap_up = UpsampleCBR(input_channels, upsample_channels, num_upsample, num_flat)
        self.paf_up = UpsampleCBR(input_channels, upsample_channels, num_upsample, num_flat)
        self.cmap_att = nn.Conv2D(upsample_channels, upsample_channels, kernel_size=3, stride=1, padding=1)
        self.paf_att = nn.Conv2D(upsample_channels, upsample_channels, kernel_size=3, stride=1, padding=1)
        self.cmap_conv = nn.Conv2D(upsample_channels, cmap_channels, kernel_size=1, stride=1, padding=0)
        self.paf_conv = nn.Conv2D(upsample_channels, paf_channels, kernel_size=1, stride=1, padding=0)    def forward(self, x):
        xc = self.cmap_up(x)
        ac =  nn.functional.sigmoid(self.cmap_att(xc))
        xp = self.paf_up(x)
        ap =  nn.functional.tanh(self.paf_att(xp))        return self.cmap_conv(xc * ac), self.paf_conv(xp * ap)

%cd /home/aistudio/work/human
!python infer.py /home/aistudio/tmp/10p.jpeg

from work.human.trt_pose_model import get_modelimport paddleimport structinput = paddle.ones((1, 3, 224, 224))
model = get_model()# print(model) #查看网络结构wgts = paddle.load("/home/aistudio/data/data127829/trt_pose.pdparams")
f = open('trt_pose.wts', 'w')
f.write('{}\n'.format(len(wgts.keys())))for k, v in wgts.items():    # print("weight key: ", k, v.shape)
    vr = v.numpy().flatten()
    f.write('{} {} '.format(k, len(vr)))    for vv in vr:
        f.write(' ')
        f.write(struct.pack('>f',float(vv)).hex())
    f.write('\n')
f.close()print("weight file created!!!")

mkdir buildcd build
cmake ..make./trt_pose -s

# 解压你所挂载的数据集在目录下!unzip -oq /home/aistudio/data/data127829/images.zip -d /home/aistudio/data
!cp /home/aistudio/data/data127829/keypoints.json /home/aistudio/data# 查看数据集的目录结构!ls /home/aistudio/data
!tree /home/aistudio/data -d

import cv2import matplotlib.pyplot as pltfrom work.animal.pre_visualize import visualize_img
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False%matplotlib inline

img = visualize_img()
plt.figure("Image") # 图像窗口名称plt.imshow(img)
plt.axis('on') # 关掉坐标轴为 offplt.title('image') # 图像题目plt.show()

%cd /home/aistudio/work/animal/trt_pose_model.py
!python train.py

%cd /home/aistudio/work/animal/
!python infer.py /home/aistudio/data/images/2007_000063.jpg

%cd /home/aistudio/work/animal/
!python infer.py /home/aistudio/data/images/ca80.jpeg