Pytorch实现单目标检测网络

单目标检测

单目标检测流程

单目标检测：判断一张图片上是否有对应目标

  • 人脸检测
  • 小黄人检测

小黄人检测图片

小黄人目标检测：需要考虑的东西

  • 数据：正负样本。
  • 二分类：判断是否包含小黄人，输出层sigmoid（数据归一化为0-1，设置一个阈值来判断是否包含小黄人）
  • 回归：画出框，输出层不需要激活(输出四个坐标值)
  • 小黄人种类：多分类问题，输出层使用softmax（每类概率和为1），小黄人也是有种类的

pytorch-Single-target-detection

数据集：包含正样本和负样本。
数据集命名：1.0.0.0.0.0.0.jpg

  • 第一个数字为序号
  • 第二个数字代表是否有小黄人，1有，0无
  • 第三到第六个数字，代表位置。
  • 第七位为种类数

实现一个单目标检测模型的流程：

  • 处理数据
  • 构建模型
  • train
  • test
  • 模型评估
  • 部署

处理数据

init初始化过程：

  • 初始化的时候，导入数据，如果是图片，一般保存图片的路径。
  • 判断是train还是test，来获取具体哪个文件夹下的。
  • 拼接文件名，获取文件。

__getitem__函数：获取某个index的图片

  • 读取self数组中的index索引图片。
  • OpenCV或者numpy读取图片，将图片进行缩放,w和h都设置为300，并对图片进行归一化。
  • 对图片进行换轴，读取到图片维度顺序为HWC，卷积操作中的tensor数据的维度是NCHW，N为Batch Size。
  • 根据图片名字，按照点，获取图片的二分类标签label，回归点位position，以及类别sort，负样本不进行后续反向传播训练。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
class MyDataset(Dataset):

    def __init__(self,root,is_train=True):
        self.dataset = []
        dir = 'train' if is_train else 'test'
        sub_dir = os.path.join(root,dir)
        img_list = os.listdir(sub_dir) # 获取文件夹下所有文件

        for i in img_list:
            img_dir = os.path.join(sub_dir,i)
            # print(img_dir)
            self.dataset.append(img_dir)
    
    def __len__(self):
        return len(self.dataset)

    def __getitem__(self,index): # 当数据被调用时就会触发该函数
        data = self.dataset[index]
        # img = cv2.imread(data) # 读到的图片是HWC ，卷积用到的是 NCHW，需要换轴 ，没有OpenCV用底下的两行代码操作代替

        # Load the image
        img = Image.open(data)
        img = img.resize((300,300))
        img = np.array(img)/255 # 读入图片并进行归一化
        
        #print(img.shape)
        # new_img = np.transpose(img,(2,0,1)) # HWC -> CHW (2,0,1) # Numpy做法 (3, 180, 298)
        new_img = torch.tensor(img).permute(2,0,1) # torch做法 torch.tensor([3, 180, 298])
        #print(new_img.shape)

        # 用点分割去取数据

        data_list = data.split('.')
        label = int(data_list[1])
        position = data_list[2:6]
        position = [int(i)/300 for i in position]
        sort = int(data_list[6])-1 # 0就变成-1了

        return np.float32(new_img),np.float32(label),np.float32(position),sort

构建网络

init 构建：

  • 具体网络设计包含：卷积，LeakyReLU，池化。卷积网络最后一层输出通道为128个通道，最后输出的卷积图像宽高为19*19
  • label的预测值：卷积128->1，对应是否包含目标
  • position的预测值：卷积128->4，对应左上角和右下角的x，y值
  • sort的预测值：卷积128->20，数据集有20类，得到最后在哪一类

forward(self,x)：

  • 实施构建网络
  • 对label，position，sort分别输出得到，并降维。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
class MyNet(nn.Module):
    def __init__(self):
        super(MyNet,self).__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(3,11,3),
            nn.LeakyReLU(),
            nn.MaxPool2d(3),
            nn.Conv2d(11,22,3),
            nn.LeakyReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(22,32,3),
            nn.LeakyReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32,64,3),
            nn.LeakyReLU(),
            nn.Conv2d(64,128,3),
            nn.LeakyReLU(),
        )

        # label的预测值
        self.label_layers = nn.Sequential(
            nn.Conv2d(128,1,19),
            nn.LeakyReLU()
        )

        # position的预测值
        self.position_layer = nn.Sequential(
            nn.Conv2d(128,4,19),
            nn.LeakyReLU()
        )

        # sort的预测值：真实的数据集有20类。
        self.sort_layer = nn.Sequential(
            nn.Conv2d(128,20,19),
            nn.LeakyReLU()
        )

    def forward(self,x):
        out = self.layers(x)

        label = self.label_layers(out)
        # 降维，从四维降为二维
        label = torch.squeeze(label,dim=2)
        label = torch.squeeze(label,dim=2)
        label = torch.squeeze(label,dim=1)

        position = self.position_layer(out)
        position = torch.squeeze(position,dim=2)
        position = torch.squeeze(position,dim=2)

        sort = self.sort_layer(out)
        sort = torch.squeeze(sort,dim=2)
        sort = torch.squeeze(sort,dim=2)

        return label,position,sort

训练逻辑

init:

  • 初始化一个summaryWriter，可视化数据的初始化。
  • loader数据集，包含train和test的，批次大小为50，打乱shuffle。
  • self设计好的网络，网络放置在GPU上，后续在GPU上训练。
  • 确定优化器为Adam
  • 二分类用BCELOGIST，里面自带着sigmod激活，网络最后输出标签不带激活函数，这里刚好可以用上。
  • 回归，MSE LOSS
  • 多分类：cross LOSS，多指交叉熵，自带softmax激活，网络最后输出类别不带激活函数，这里刚好可以用上。

call函数：

  • 设定训练epoch，默认是1000
  • 读取数据集和标签img, label, position, sort，将读取到的数据放到GPU上，后续在GPU上训练。
  • 将图片放入卷积网络中得到out_label,out_position,out_sort，预测值。
  • 将三个值分别和标签进行对比，计算各个的loss，相加得到总loss，这里各个loss的权重一样，可以独立设置参数。
  • 清空梯度，反向传播，更新权重
  • 每隔十次打印一下训练损失。
  • 每个epoch，保存一个训练权重。
  • test逻辑和训练逻辑差不多，不过不需要进行梯度更新。将图片放入网络，得到结果后和标签进行损失计算，计算标签准确率和平均值。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
class Train:
    def __init__(self,root,weight_path):
        self.summaryWriter = SummaryWriter('logs')
        self.train_dataset = MyDataset(root=root,is_train= True)
        self.test_dataset = MyDataset(root=root,is_train= False)
        # 挺好，现在知道这玩意是在干嘛了
        self.train_dataLoader = DataLoader(self.train_dataset,batch_size=50,shuffle=True)
        self.test_dataLoader = DataLoader(self.test_dataset,batch_size=50,shuffle=True)

        self.net = MyNet().to(DEVICE) # 放到gpu上

        if os.path.exists(weight_path):
            self.net.load_state_dict(torch.load(weight_path))

        self.opt = optim.Adam(self.net.parameters())

        # 分类，回归，都有损失，怎么计算？
        
        self.label_loss_fun = nn.BCEWithLogitsLoss()  # 二分类用BCELOGIST，里面自带着sigmod激活
        self.position_loss_fun = nn.MSELoss() # 回归，MSE LOSS
        self.sort_loss_fun = nn.CrossEntropyLoss() # 多分类：cross LOSS，多指交叉熵，自带softmax激活

        self.train = True
        self.test = True
        
    def __call__(self):
        index1,index2 = 0,0
        for epoch in range(1000): # 训练1000个epoch
            if self.train:
                for i, (img,label,position,sort) in enumerate(self.train_dataLoader):
                    self.net.train()
                    
                    img, label, position, sort = img.to(DEVICE), label.to(DEVICE), position.to(DEVICE), sort.to(DEVICE) # 数据放到GPU上
                    # print(img.shape)
                    # print(label.shape)
                    out_label,out_position,out_sort = self.net(img)
                    # print(out_label,out_position,out_sort) 
                    # print(out_label.shape)
                    
                    label_loss = self.label_loss_fun(out_label,label)
                    # print(label_loss)

                    position_loss=self.position_loss_fun(out_position,position)
                    # print(position.shape)
                    # print(out_position.shape)
                    # print(position_loss)
                    
                    sort = sort[torch.where(sort >= 0)]
                    out_sort = out_sort[torch.where(sort >= 0)]
                    sort_loss = self.sort_loss_fun(out_sort,sort)
                    # print(sort_loss)

                    train_loss = label_loss + position_loss + sort_loss
                    
                    self.opt.zero_grad()
                    train_loss.backward()
                    self.opt.step()

                    if i%10 ==0 :
                        print(f'train_loss{i}===>',train_loss.item())
                        self.summaryWriter.add_scalar('train_loss',train_loss,index1)
                        index1 +=1
                    
                data_time = str(datetime.datetime.now()).replace(' ', '-').replace(':','_').replace('·','_')
                save_dir = 'param'
                if not os.path.exists(save_dir):
                    os.makedirs(save_dir)

                # 保存权重文件
                torch.save(self.net.state_dict(), f'{save_dir}/{data_time}-{epoch}.pt')
            
            if self.test:
                sum_sort_acc,sum_label_acc = 0,0
                for i, (img,label,position,sort) in enumerate(self.test_dataLoader):
                    self.net.train()
                    
                    img, label, position, sort = img.to(DEVICE), label.to(DEVICE), position.to(DEVICE), sort.to(DEVICE) # 数据放到GPU上
                    
                    out_label,out_position,out_sort = self.net(img)
                    
                    label_loss = self.label_loss_fun(out_label,label)

                    position_loss=self.position_loss_fun(out_position,position)
                    
                    # print(position_loss)
                    
                    sort = sort[torch.where(sort>=0)]
                    out_sort = out_sort[torch.where(sort >= 0)]
                    sort_loss = self.sort_loss_fun(out_sort,sort)
                    # print(sort_loss)

                    test_loss = label_loss + position_loss + sort_loss

                    out_label = torch.tensor(torch.sigmoid(out_label))
                    out_label[torch.where(out_label>=0.5)] = 1
                    out_label[torch.where(out_label<0.5)] = 0

                    label_acc = torch.mean(torch.eq(out_label,label).float())
                    sum_label_acc += label_acc 

                    # 求准确率
                    # out_sort = torch.argmax(torch.softmax(out_sort,dim=1))
                    if out_sort.numel() > 0:
                        out_sort = torch.argmax(torch.softmax(out_sort, dim=1))
                        out_sort = out_sort.to(sort.device)  # Move out_sort to the same device as sort
                    else:
                        out_sort = torch.tensor([], device=sort.device)  # Or handle the empty case 

                    sort_acc = torch.mean(torch.eq(sort,out_sort).float())
                    sum_sort_acc += sort_acc
                    
                    if i%10 ==0 :
                        print(f'test_loss{i}===>',test_loss.item())
                        self.summaryWriter.add_scalar('test_loss',test_loss,index2)
                        index2 +=1
                    
                avg_sort_acc = sum_sort_acc/i
                print(f'avg_sort_acc {epoch}====>',avg_sort_acc)
                self.summaryWriter.add_scalar('avg_sort_acc',avg_sort_acc,epoch)

                avg_label_acc = sum_label_acc/i
                print(f'avg_label_acc {epoch}====>',avg_label_acc)
                self.summaryWriter.add_scalar('avg_label_acc',avg_label_acc,epoch)

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

预测流程：

  • 读取test文件夹，依次读取img
  • 读取标签值和坐标值，输出在图片上。
  • 网络预测标签和坐标值，也输出在图片上，可以直观的对比网络预测值和标签值的差异。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
if __name__ == '__main__':
    img_name = os.listdir('/home/liuchang/codetest/PythonCode/YellowPersonDetect/yellow_data/test') # 获取文件路径
    for i in img_name:
        img_dir = os.path.join('/home/liuchang/codetest/PythonCode/YellowPersonDetect/yellow_data/test',i) # 获取img

        img = cv2.imread(img_dir)

        position = i.split('.')[2:6] # 读取位置坐标
        position = [int(j) for j in position]
        sort =  i.split('.')[6]
        cv2.rectangle(img,(position[0],position[1]),(position[3],position[4]),(0,255,0),thickness=2) # 画label框
        cv2.putText(img,sort,(position[0],position[1]-3),cv2.FONT_HERSHEY_SIMPLEX,2,1,(0,255,0),thickness=2) #左上角的点上写出类型

        model = MyNet()
        model.load_state_dict(torch.load('param/2024-04-10-14_37_43.382495-50.pt'))
        new_img = torch.tensor(img).permute(2,0,1)
        torch.unsqueeze(new_img,dim=0)/255 # 传递一个维度，归一化

        label_out,out_position,out_sort = model(new_img)
        label_out = torch.sigmoid(label_out) # 网络模型没有做归一化，这里就要做归一化
        out_sort = torch.argmax(torch.softmax(out_sort,dim=1))

        out_position = out_position[0]*300
        out_position = [int(i) for i in out_position]
        if label_out.item()>0.5 :
            cv2.rectangle(img,(out_position[0],out_position[1]),(out_position[3],out_position[4]),(255,0,0),thickness=2) # 画label框
            cv2.putText(img,str(out_sort.item()),(out_position[0],out_position[1]-3),cv2.FONT_HERSHEY_SIMPLEX,2,1,(255,0,0),thickness=2) #左上角的点上写出类型
        
        cv2.imshow('img',img)
        cv2.waitKey(500)
        cv2.destroyAllWindows() # 展示图片后关闭图片

完整代码：https://github.com/cauccliu/SingleObjectDetect

参考列表：

人工智能
#Python #CV
Pytorch实现单目标检测网络
https://cauccliu.github.io/2024/04/17/YelloPersonDetect/
Author
Liuchang
Posted on
April 17, 2024
Licensed under