关键点检测——直接回归法

一、数据集格式

 二、解析xml文件，生成data_center.txt

from PIL import Image
import math,os
from xml.etree import ElementTree as ET


def keep_image_size_open(path, size=(256, 256)):
    img = Image.open(path)
    temp = max(img.size)
    mask = Image.new('RGB', (temp, temp), (0, 0, 0))
    mask.paste(img, (0, 0))
    mask = mask.resize(size)
    return mask


def make_data_center_txt(xml_dir):
    with open('data_center.txt', 'a') as f:
        f.truncate(0)
        path=r'data/images'
        xml_names = os.listdir(xml_dir)
        for xml in xml_names:
            xml_path = os.path.join(xml_dir, xml)
            in_file = open(xml_path)
            tree = ET.parse(in_file)
            root = tree.getroot()
            image_path = root.find('path')
            polygon = root.find('outputs/object/item/polygon')
            data = []
            c_data = []
            data_str = ''
            print(xml)
            for i in polygon:
                data.append(int(i.text))
                data_str = data_str + ' ' + str(i.text)
            for i in range(0, len(data), 2):
                c_data.append((data[i], data[i + 1]))
            data_str = os.path.join(path,image_path.text.split('\\')[-1]) +data_str
            f.write(data_str + '\n')


if __name__ == '__main__':
    make_data_center_txt('data/xml')

 三、加载数据集

import torch
from torch.utils.data import Dataset
from torchvision import transforms
from PIL import Image

tf = transforms.Compose([  #标准化处理
    transforms.ToTensor()
])

class MyDataset(Dataset):
    def __init__(self,root): #传入路径
        f=open(root,'r')
        self.dataset=f.readlines() #读所有行
    def __len__(self):
        return len(self.dataset) #返回数据集长度
    def __getitem__(self, index):
        data=self.dataset[index] #取当前数据
        img_path=data.split(' ')[0] #以空格划分，并取出文件名，即data/images\0.png
        img_data=Image.open(img_path) #打开图片
        # points = data.split(' ')[1:-2]  # 取出后面5个点的x,y坐标，-2是取不到的
        points=data.split(' ')[1:] #取出后面5个点的x,y坐标
        # print(img_data, points)
        points = [int(points[0])/774, int(points[1])/434, int(points[2])/774, int(points[3])/434, int(points[4])/774, int(points[5])/434]
        # points=[int(i)/100 for i in points] #图像宽高为100，int(i)/100进行归一化
        # print(img_data, points)
        return tf(img_data),torch.Tensor(points) #将img_data标准化，将points转化为tensor格式


if __name__ == '__main__':
    data=MyDataset('data_center.txt')
    for i in data:
        print(i[0].shape)
        print(i[1].shape)

四、构建网络

import torch
from torchvision import models
from torch import nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layer=nn.Sequential( #用resnet50模型
            models.resnet50(pretrained=True)
        )
        #全连接层的输出要改为自己对应的输出，将1000分类通过全连接层变为6分类
        self.out=nn.Linear(1000,6)
    def forward(self,x):
        return self.out(self.layer(x)) #将输入x经过resnet50以及全连接层Linear

if __name__ == '__main__':
    net=Net()
    x=torch.randn(1,3,100,100)
    print(net(x).shape)

五、开始训练

import os
from torch import nn,optim
import torch
from dataset import *
from net import *
from torch.utils.data import DataLoader


if __name__ == '__main__':
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    net=Net().to(device) #实例化网络并指认到设备上
    weights='params/net.pth'
    if os.path.exists(weights): #如果有初始权值就加载
        net.load_state_dict(torch.load(weights)) #加载权重
        print('loading successfully')
    opt=optim.Adam(net.parameters()) #指定优化器并传入参数
    loss_fun=nn.MSELoss() #定义损失函数
    dataset=MyDataset('data_center.txt') #实例化数据集
    data_loader=DataLoader(dataset,batch_size=2,shuffle=True) #加载数据集
    epoch = 1
    while True:
        for i,(image,label) in enumerate(data_loader): #用枚举的方式遍历数据集
            image,label=image.to(device),label.to(device) #将图片和标签指认到设备上
            # print(image.shape, label.shape)
            out=net(image) #将图片输入网络
            train_loss=loss_fun(out,label) #预测值和真是标签做损失

            print(f'{epoch}-{i}-train_loss:{train_loss.item()}') #打印当前轮次当前批次的训练损失

            opt.zero_grad() #梯度清零
            train_loss.backward() #反向传播
            opt.step() #更新梯度
        if epoch % 10 == 0: #每10轮保存一次权重
            torch.save(net.state_dict(),f'params/net.pth') #保存参数
            print('save successfully')
        epoch += 1

六、利用训练好的权重进行预测

import os
import torch
from PIL import Image,ImageDraw
from dataset import *
from net import *    #import * 代表导入所有

path='test_image'
net=Net() #实例化网络
net.load_state_dict(torch.load('params/net.pth')) #加载训练好的权重
net.eval() #测试模式
for i in os.listdir(path):
    img=Image.open(os.path.join(path,i))
    draw=ImageDraw.Draw(img) #创建画板
    img_data=tf(img) 
    img_data=torch.unsqueeze(img_data,dim=0)
    out=net(img_data)
    # print(out, out.shape)
    out=(out[0]).tolist() #取第0个，并由tenser转化成列表形式
    out = [out[0]*774,out[1]*434,out[2]*774,out[3]*434,out[4]*774,out[5]*434]

    # print(out)
    for j in range(0,len(out),2):
        draw.ellipse((out[j]-2,out[j+1]-2,out[j]+2,out[j+1]+2),(255,0,0)) #画半径为2的圆
    img.show()

七、制作数据集 

精灵标注助手->选择多边形框标注->标注完一张Ctrl+S保存->导出XML格式

reference

>>>>>来自B站大佬

【深度学习关键点回归(直接回归法&heatmap热力图法)】 https://www.bilibili.com/video/BV1sS4y197J1/?p=2&share_source=copy_web&vd_source=95705b32f23f70b32dfa1721628d5874