pytorch 实现目标检测（一）（小黄人检测训练）

一、项目介绍

该示例首先有两部分数据，第一部分准备小黄人数据集，第二部分准备自然图像数据集。将小黄人图片粘贴到自然图像上面，然后检测图片中的小黄人，下面是粘贴好的图片。实验中可以用一个较小的数据集进行训练，准备5000张自然图片，小黄人可以下载10 - 20 多张不同的，将其随机粘贴到5000张自然图像上即可。

二、数据预处理

import numpy as np
from PIL import Image
import os
# 对收集的5000张自然图像进行处理，因为利用爬虫爬的图片大小，图片的模式都不相同，所以为了数据的统一先进行处理。
def convertImage():
    listpath = os.listdir("path") # 这里path 表示你大自然图片的数据集路径。
    for path in listpath:
        img = Image.open(os.path.join("path",path))
        img = img.convert("RGB")   # 转化为RGB模式
        # 尺寸设置为224*224，采用细节增强缩放  ，
        img = img.resize((224,224),Image.ANTIALIAS) 
        img.save("")  # 保存处理好的图片

def createDataset(dirimage):
    listpath = os.listdir(dirimage)
    for index , path in enumerate(listpath): # enumerate 这个函数的作用就是增加一个索引值（可以理解为加序号）
        img = Image.open(os.path.join(dirimage,path))
        “”“ 
        这里将数据集分为训练集4000张（2000张作为正样本，也就是贴上小黄人的样本，2000张作为负样），
        测试集1000张（500张作为正样本，500张作为负样本）。0-1999, 4000-4499这两个区间上的图片粘贴小黄人，其他的不粘贴。
        ”“”
        if index < 2000 or (index >= 4000 and index < 4500):
            """随机取出一张小黄人图片，这里还可以用np.random.choice(img_yellow)，img_yellow 是一个列表，这是从列表中随机取一个值的方式"""
            minions=Image.open("yellow/{}.png".format(np.random.randint(1,21)))
            # 缩放
            h = w = np.random.randint(64,180) 
            # 将小黄人的尺寸随机缩放成64 - 180 的大小
            minions = minions.resize((h,w),Image.ANTIALIAS)
            # 旋转
            minions = minions.rotate(np.random.randint(-30,30))
            # 翻转，镜像翻转
            minions = minions.transpose(Image.FLIP_LEFT_RIGHT) if np.random.randint(0,2) == 1 else minions
            x,y = np.random.randint(0, 224 - w) , np.random.randint(0, 224 - h)
            # 掩码
            r,g,b,a = minions.split() # 分离通道小黄人是RGBA 的格式
            # mask 表示图像的一部分，或者说你感兴趣的部分。这里粘贴小黄人的a通道
            img.paste(minions,(x,y), mask=a)  

            # print(x,y)
            if not os.path.isdir("datasets"): # 判断文件夹是否存在
                os.mkdir("datasets")   # 创建文件夹
            img.save("datasets/{}.{}.{}.{}.{}.{}.jpg".format(index,x,y,x+w,y+h,1))
        else:
            img.save("datasets/{}.0.0.0.0.0.jpg".format(index)


if __name__ == '__main__':
    convertImage()
    createDataset(r"图片路径")

三、制作数据集

制作数据集基本的思路（结构），加载一个类，定义三个模块。具体看下面代码：

# 数据集制作的基本框架
from torch.utils.data import Dataset
“”“`Dataset`代表Dataset的抽象类，这个模块请阅读torch底层代码“”“

class MyDateSet(Dataset):
    # 该模块的作用是加载数据集 
    def __init__(self):
        super(MyDateSet, self).__init__()
        pass
    # 获取数据集的大小
    def __len__(self):
        pass
    # 根据索引取出数据，返回数据和数据的标签。
    def __getitem__(self, index):
        pass

下面利用这个三个模块制作数据集。

import torch
from torch.utils.data import Dataset,DataLoader
import os
from PIL import Image
import numpy as np
import torchvision.transforms as trans


class MyDataset(Dataset):
    # 计算mean 和std 的方法请看我的博客《pytorch 猫狗数据集识别（一）》这里面有详细的计算过程
    mean = torch.tensor([0.5708, 0.5661, 0.5395])
    std = torch.tensor([0.3128, 0.2978, 0.3172])
    
    def __init__(self,root=None,train=True,transforms=None):
        self.path = root
        self.transforms = transforms
        self.dataset = os.listdir(self.path)
        # index(".")查找第一个"."的位置，然后x[:x.index(".")]取出该值：
        “”“                                    序号 x1 y1 x2 y2 标签（1表示粘贴了小黄人）
        可能这样说不是很清楚，举个例子：图片的标签是a['12.35.67.68.97.1.jpg']
        a[0] = '12.35.67.68.97.1.jpg'
        a[0].index(".") == 2
        a[0][:a[0].index(".")] 即a[0][:2]==12     
        所以可以看出这个代码是为了取出图片的序号，也就是图片的第一项。
        ”“”
        self.dataset.sort(key=lambda x: int(x[:x.index(".")])) #根据图片的序号进行排序
        if train:
             # train == True 0-3999为训练集 这里2000张正样本，2000张负样本
            self.dataset = self.dataset[:4000]  
        else:
            # train ==False 4000-4999为测试集，这里500张正样本，500张负样本
            self.dataset = self.dataset[4000:5000] 
    # 获取数集的大小
    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, index):
        imgpath = self.dataset[index] # 跟据索引获取图片
        img = Image.open(os.path.join(self.path, imgpath))
        # transforms=trans.Compose([trans.ToTensor(), trans.Normalize(MyDataset.mean, MyDataset.std)])
        data = self.transforms(img) 
        labels = imgpath.split(".")
        # 这里取的是坐标，必须要将坐标转换为array，这样才能进行计算，并且除以224对其进行归一化处理
        axes = np.array(labels[1:5],dtype=np.float32) / 224 
        category = np.array(labels[5:6], dtype=np.float32) # 这里取的是标签，1 和 0 
        # 拼接列表，这里有先后顺序
        target = np.concatenate((axes,category))
        return data, target

四、构建网络模型

# 计算特征图大小的计算公式：
“”“
out_size = (输入图片的大小 - 卷积核大小 + 2 * padding) / 步长  + 1 
”“”

import torch.nn as nn
# import numpy as np
import torch


class MyNetWork(nn.Module):
    def __init__(self):
        super(MyNetWork, self).__init__()
        self.convlution_layer =nn.Sequential(
            """ 
            1、这里用卷积 激活 卷积 激活  池化，用了两层卷积目的是用小的卷积核代替大的卷积核
            2、还要注意一点，卷积的时候一般在最前面要用最大池化，网络后面用平均池化
            3、nn.Conv2d(输入通道，输出通道，卷积核大小，步长，padding)
            4、nn.MaxPool2d(卷积核大小，步长)
            5、nn.ReLu(inplace=True) # 
            利用in-place计算可以节省内（显）存，同时还可以省去反复申请和释放内存的时间。但是会对原变量覆盖。也即对从上层网络nn.Conv2d中传递下来的tensor直接进行修改，这样能够节省运算内存，不用多存储其他变量。具体的看这篇博客：https://blog.csdn.net/AugustMe/article/details/92589979
            """
            
            nn.Conv2d(3, 16, 3, 1),  # (224 - 3 + 2 * 0)/1+1 =222
            nn.ReLU(inplace=True),
            nn.Conv2d(16, 32, 3, 1), # (222 - 3 + 2*0)/1 +1 = 220
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),  # (220 - 2 + 2*0) / 2 +1 = 110
            nn.Conv2d(32, 128, 3, 1), # (110 -3 + 2*0) /1 +1 =108
            nn.ReLU(inplace=True),
            # nn.Conv2d(64, 128, 3, 1), # (108 - 3 + 2*0 )/1 + 1 =106
            # nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),   # (108 -2 + 2*0)/2 + 1 = 54
            nn.Conv2d(128, 256, 3, 1), # (54 - 3 + 2*0)/1 + 1 = 52
            nn.ReLU(inplace=True),
            # nn.Conv2d(256, 64, 3, 1), # (51 - 3 + 2*0)/1 + 1 = 49
            # nn.ReLU(inplace=True),
            nn.AvgPool2d(2, 2),    # (52 -2 + 2*0) /2 +1 =26
            nn.Conv2d(256, 64, 3, 1), # (26 - 3)/1 +1= 24
            nn.ReLU(inplace=True),
            nn.AvgPool2d(2, 2),  # (24 -2 )/2 +1 = 12
            nn.Conv2d(64, 32, 3, 1) # (12 - 3)/1 +1 = 10
        )
       # 利用全连接网络对特征进行拼接。其实是为了将特征组合后输出五个结果（坐标 + 标签）
        self.MLP_layer = nn.Sequential(
            nn.Linear(32*10*10, 128), # 输出了32个特征图，且特征图大小为10*10 
            nn.ReLU(),
            nn.Linear(128, 5)
        )
        # 或者利用全卷积对特征进行拼接，用卷积核大小为1*1的进行卷积
	  #self.cnn_layer2 = nn.Sequential(
      #      nn.Conv2d(32,128,10,1),  # 这里用10*10 的卷积核，因为特征图是10*10大小的，进行卷积后就得到大小为1*1的特征图
      #      nn.ReLU(True),
      #      nn.Conv2d(128,5,1,1)   # 输出5个结果，用1*1大小的卷积核进行卷积
        )
    def forward(self, x):
        input = self.convlution_layer(x)
        x_out = torch.reshape(input, shape=(-1, 32*10*10))
        out = self.MLP_layer(x_out)
        category = torch.sigmoid(out[:, 4]) # 最后一个为标签0, 1，sigmoid表示概率，一般用于二分类，因为它的输出值只会偏向两边，softmax常用于多分类。（这里用sigmoid作分类）
        axes = torch.relu(out[:, :4]) #前四个为坐标 （这里作回归，因为坐标是正数，所以用relu来解决）
        return axes, category

五、训练网络

import torch
from torch.utils.data import DataLoader
import torch.nn as nn
import torchvision.transforms as trans
from MyNet import MyNet
from MyData import MyDataset
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
import numpy as np
import os

class Trainer:

    def __init__(self):
        transforms = trans.Compose([ 
            trans.ToTensor(),
            trans.Normalize(MyDataset.mean,MyDataset.std)
        ])
        self.train_dataset = MyDataset(root=r"/home/gwp/PycharmProjects/ObjectDetection/train_Image", train=True, transforms=transforms)
        self.test_dataset = MyDataset(root=r"/home/gwp/PycharmProjects/ObjectDetection/train_Image", train=False, transforms=transforms)
        self.net = MyNet()
        # 回归问题最常用的损失函数是均方误差MSELoss，定义如下：
        self.offset_lossfunc = nn.MSELoss().cuda()   

         # BCELoss(Binary Cross Entropy Loss)，就是交叉熵应用于二分类时候的特殊形式，一般都和sigmoid一起用.
        self.category_lossfunc = nn.BCELoss().cuda()
        self.optimier = torch.optim.Adam((self.net.parameters()))

    def train(self):
       # 这里需要注意下，如果以后在训练的过程中出现了训练中断了，可以加载保存的网络模型继续接着训练。
        if os.path.exists("models/net2.pth"):
            self.net = torch.load("models/net2.pth")
            print("exists")

        trainloader = DataLoader(dataset=self.train_dataset, batch_size=108, shuffle=True)
        losses = []
        for i in range(20):
            print("epochs:{}".format(i))
            for j, (x , y) in enumerate(trainloader):
                if torch.cuda.is_available():
                    x = x.cuda()
                    y = y.cuda()
                category,axes = self.net(x)  # 前向传播输出坐标和标签
                # (106,)   (106,)
                loss1 = self.category_lossfunc(category, y[:,4])
                loss2 = self.offset_lossfunc(axes, y[:,0:4])# (106,4)
                loss = loss1 + loss2
                if j % 5 == 0:
                    losses.append(loss.float())
                    print("{}/{},loss:{}".format(j, len(trainloader), loss.float()))
                    plt.clf()
                    plt.plot(losses)
                    plt.pause(0.1)

                self.optimier.zero_grad()
                loss.backward()
                self.optimier.step()
                del x,y, category,axes,loss1,loss2,loss
            torch.save(self.net, "models/net2.pth")