手写--单目标检测

做3件事：

1. 读取数据
2. 搭建网络
3. 训练代码

在目标检测任务中通常有3个损失需要进行训练：1.图片是否存在检测目标、2.目标左上角与右下角坐标位置、3.目标为什么类别

读取数据代码：

from torch.utils.data import Dataset
import os
import cv2
import numpy as np
import torch


class MyDataset(Dataset):
    # 导入数据 保存路径即可

    def __init__(self, is_train=True):

        self.dataset = []
        self.train = "data/images/train/"
        self.test = "data/images/test/"

        self.lab = []
        self.train_lab = "data/labels/train/"
        self.test_lab = "data/labels/test/"

        if is_train:
            # 读取图像
            path_img = os.listdir(self.train)
            # print(path_img)
            for i in path_img:
                self.dataset.append(os.path.join(self.train, i))
            # 读取标签
            path_lab = os.listdir(self.train_lab)
            for j in path_lab:
                self.lab.append((os.path.join(self.train_lab, j)))

        else:
            # 读取图像
            path_img = os.listdir(self.test)
            # print(path_img)
            for i in path_img:
                self.dataset.append(os.path.join(self.test, i))
            # 读取标签
            path_lab = os.listdir(self.test_lab)
            for j in path_lab:
                self.lab.append((os.path.join(self.test_lab, j)))

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, index):
        data = self.dataset[index]
        img = cv2.imread(data)
        img = img / 255
        # img = cv2.resize(img, (300, 300))

        # cv2的图像：           H-高  W-宽  C-通道
        # 卷积需要的图像：N-批次、C-通道、H-高、W-宽、

        # 根据索引 交换位置：cv2 的索引是正常的 0,1,2 现在需要换位置

        # 方法一 np的方法
        # new_img = np.transpose(img, (2, 0, 1))

        # 方法2 torch的方法
        new_img = torch.tensor(img).permute(2, 0, 1)
        # print(new_img)

        # 读取标签
        lab = self.lab[index]
        fo = open(lab, "r")

        label, str_1, str_2, str_3, str_4, str_class = fo.readline().split(" ")

        return np.float32(new_img), np.float32(label),\
               np.float32([int(str_1) / 300, int(str_2) / 300, int(str_3) / 300, int(str_4) / 300]),\
               int(str_class)


if __name__ == '__main__':
    path = MyDataset()
    for i in path:
        print(i)

搭建网络

from torch import nn
import torch


class Mynet(nn.Module):
    def __init__(self):
        super(Mynet, self).__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(3, 16, 3),
            nn.ReLU(),
            nn.MaxPool2d(3),

            nn.Conv2d(16, 22, 3),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(22, 32, 5),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(32, 64, 5),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(64, 82, 3),
            nn.ReLU(),

            nn.Conv2d(82, 128, 3),
            nn.ReLU(),

            nn.Conv2d(128, 25, 3),
            nn.ReLU(),
        )


        # 是否为背景
        self.label_layer = nn.Sequential(
            nn.Conv2d(25, 1, 3),
            nn.ReLU(),
        )

        # x1 y1 x2 y2 坐标位置
        self.position_layer = nn.Sequential(
            nn.Conv2d(25, 4, 3),
            nn.ReLU(),
        )
        # 多类别
        self.sort_layer = nn.Sequential(
            nn.Conv2d(25, 20, 3),
            nn.ReLU(),
        )

    def forward(self, x):
        out = self.layers(x)
        # 当dim不设值时，去掉输入的tensor的所有维度为1的维度;
        # 当dim为某一整数（0 <= dim < input.dim()）时，判断dim维的维度是否为1，若是则去掉，否则不变。
        # 另外，当input是一维的时候，squeeze不变

        # dim = 3 是指 torch.Size[1,1,1,1] 从0开始 左向右数 到3，也就是第4个
        label = self.label_layer(out)
        label = torch.squeeze(label, dim=3)
        label = torch.squeeze(label, dim=2)
        label = torch.squeeze(label, dim=1)

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

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

        return label, position, sort


if __name__ == '__main__':
    net = Mynet()
    x = torch.randn(3, 3, 300, 300)

    print(net(x)[0].shape)
    print(net(x)[1].shape)
    print(net(x)[2].shape)

训练代码

from net import Mynet
from data import MyDataset
from torch.utils.data import DataLoader
# from torch.utils.tensorboard import SummaryWriter
from torch import nn, optim
import torch
import datetime
import os

DEVICE = 'cuda'


class Train:
    def __init__(self, weight_path):
        # self.summaryWriter = SummaryWriter('logs')

        self.train_dataset = MyDataset(is_train=True)
        # self.test_dataset = MyDataset(is_train=False)

        # 加载数据
        # 打乱shuffle
        self.train_dataLoader = DataLoader(self.train_dataset, batch_size=50, shuffle=True)
        # self.test_dataloader = DataLoader(self.test_dataset, batch_size=4, shuffle=True)

        # 加载模型
        self.net = Mynet().to(DEVICE)

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

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

        # 2分类
        self.label_loss_fun = nn.BCEWithLogitsLoss()  # 自带激活
        # 回归
        self.positon_loss_fun = nn.MSELoss()
        # 多分类
        self.sort_loss_fun = nn.CrossEntropyLoss()  # 自带

        # self.train = True
        # self.test = True

    def __call__(self):
        for epoch in range(1000):
            for i, (img, label, position, sort) in enumerate(self.train_dataLoader):

                img, label, position, sort = img.to(DEVICE), label.to(DEVICE), position.to(DEVICE), sort.to(DEVICE)
                # 对3个进行优化器求解
                out_label, out_osition, out_sort = self.net(img)

                label_loss = self.label_loss_fun(out_label, label)
                osition_loss = self.positon_loss_fun(out_osition, position)
                sort_loss = self.sort_loss_fun(out_sort, sort)

                train_loss = 0.2 * label_loss + osition_loss * 0.6 + 0.2 * sort_loss

                self.opt.zero_grad()
                train_loss.backward()
                self.opt.step()

                if i % 10 == 0:
                    print(f'第{i}次  train_loss = {train_loss.item()}')
            print(f"第{epoch}轮")
            date_time = str(datetime.datetime.now()).replace(' ', '-').replace(':', '_').replace('.', '_')
            torch.save(self.net.state_dict(), f'p/{date_time}-{epoch}.pt')


if __name__ == '__main__':
    train = Train("2023-03-19-20_54_33_272814-17.pt")
    train()

测试代码

from net import Mynet
import cv2
import torch


def cv_show(neme, img):
    cv2.imshow(neme, img)  # 必要参数：名字和变量名
    cv2.waitKey(0)  # 括号中0=任意键终止，单位为毫秒级别
    cv2.destroyAllWindows()  # 关闭所有窗口--图片


model = Mynet()
model.load_state_dict(torch.load('p/1.pt', map_location=torch.device('cpu')))

img = cv2.imread("1.jpg")
img = cv2.resize(img, (300, 300))
img_data = torch.tensor(img).permute(2, 0, 1)
img_data = torch.unsqueeze(img_data, dim=0) / 255
rst = model(img_data)
label = torch.sigmoid(rst[0])
sort = torch.softmax(rst[2], dim=1)


rst_1 = rst[1] * 300
x1, y1, x2, y2 = rst_1.cpu().detach().numpy()[0]


# 绘制矩形 参数：图像、左上坐标、右下坐标、颜色(蓝B,绿G,红R)、粗细
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 3)
cv_show('neme', img)

参考：github源码地址：https://github.com/qfs1980398040/pytorch-Single-target-detection-Minions
参考：原始数据集链接：https://pan.baidu.com/s/1AVifpTWr29BEfQYGhuEd_w
提取码：p5rj