【程序】目标检测 train test net datasets 手敲版

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


class MyDataset(Dataset):
    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)

        self.dataset.sort(key=lambda x: int(x[:x.index(".")]))

        if train:
            self.dataset = self.dataset[:4000]
        else:
            self.dataset = self.dataset[:1000]

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

    def __getitem__(self, index):
        imgpath = self.dataset[index]
        img = Image.open(os.path.join(self.path, imgpath))
        data = self.transforms(img)
        labels = imgpath.split(".")
        axes = np.array(labels[1:5], dtype=np.float32) / 224
        category = np.array(labels[5:6], dtype=np.float32)
        target = np.concatenate((axes, category), axis=None)
        return data, target


class MyNetWork(nn.Module):
    def __init__(self):
        super(MyNetWork, self).__init__()
        self.convlution_layer1 = nn.Sequential(
            nn.Conv2d(3, 16, 3, 1),  # 224 -> 222
            nn.ReLU(inplace=True),  #
            nn.Conv2d(16, 32, 3, 1),  # 222 -> 220
            nn.ReLU(inplace=True),

            nn.MaxPool2d(2, 2),  # 220-> 110
            nn.Conv2d(32, 128, 3, 1),  # 110->108
            nn.ReLU(inplace=True),

            nn.AvgPool2d(2, 2),  # 108->54
            nn.Conv2d(128, 256, 3, 1),  # 54-> 52
            nn.ReLU(inplace=True),
            nn.AvgPool2d(2, 2),  # 52->26
            nn.Conv2d(256, 64, 3, 1),  # 26-> 24
            nn.ReLU(inplace=True),
            nn.AvgPool2d(2, 2),  # 24->12
            nn.Conv2d(64, 32, 3, 1)  # 12->10
        )

        self.MLP_layer = nn.Sequential(
            nn.Linear(32 * 10 * 10, 128),
            nn.ReLU(),
            nn.Linear(128, 5)
        )

    def forward(self, x):
        input = self.convlution_layer1(x)
        x_out = torch.reshape(input, shape=(-1, 32 * 10 * 10))
        out = self.MLP_layer(x_out)
        category = torch.sigmoid(out[:, 4])
        axes = torch.relu(out[:, :4])
        return axes, category


def countIOU(recA, recB):
    xA = max(recA[0], recB[0])
    yA = max(recA[1], recB[1])
    xB = min(recA[2], recB[2])
    yB = min(recA[3], recB[3])

    interarea = max(0, xB - xA) * max(0, yB - yA)

    recA_area = (recA[2] - recA[0]) * (recA[3] - recA[1])
    recB_area = (recB[2] - recB[0]) * (recB[3] - recB[1])

    return interarea / float(recA_area + recB_area - interarea)


class Trainer:
    def __init__(self):
        transforms1 = trans.Compose([
            trans.ToTensor(),
            trans.Normalize(MyDataset.mean, MyDataset.std)
        ])
        self.train_dataset = MyDataset(root=r"../train_picture", train=True, transforms=transforms1)
        self.test_dataset = MyDataset(root=r"../test_picture", train=False, transforms=transforms1)

        self.net = MyNetWork().cuda()
        self.offset_lossfunc = nn.MSELoss().cuda()

        self.category_lossfunc = nn.BCELoss().cuda()
        self.optimier = torch.optim.Adam(self.net.parameters(), lr=1e-5)

    def train(self):
        trainloader = DataLoader(dataset=self.train_dataset, batch_size=50, shuffle=True)
        print(trainloader)
        losses = []
        for i in range(50):
            print("epochs:{}".format(i))
            for j, (x, y) in enumerate(trainloader):
                print(j)
                if torch.cuda.is_available():
                    x = x.cuda()
                    y = y.cuda()
                axes, category = self.net(x)

                loss1 = self.offset_lossfunc(axes, y[:, 0:4])
                loss2 = self.category_lossfunc(category, y[:, 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()
                # print(self.optimier.state)
                loss.backward()
                self.optimier.step()
                del loss, loss1, loss2, x, y, category, axes
        torch.save(self.net, "../models/net33.pth")

    def test(self):
        testloader = DataLoader(dataset=self.test_dataset, batch_size=50, shuffle=False)
        self.net = torch.load("../models/net33.pth")
        total = 0
        nnn = 0
        temp = 0
        total_iou = 0
        for x, y in testloader:
            print(temp)
            temp += 1
            x, y = x.cuda(), y.cuda()
            axes, category = self.net(x)
            total+=(category.round()==y[:,4]).sum()
            index= category.round() == 1
            target = y[index]
            x = (x[index].cpu() * MyDataset.std.reshape(-1, 3, 1, 1) + MyDataset.mean.reshape(-1, 3, 1, 1))

            for j, i in enumerate(axes[index]):
                boxes = (i.data.cpu().numpy() * 224).astype(np.int32)
                target_box = (target[j, 0:4].data.cpu().numpy() * 224).astype(np.int32)
                img = trans.ToPILImage()(x[j])
                total_iou += countIOU(boxes, target_box)

                plt.clf()
                plt.axis("off")
                draw = ImageDraw.Draw(img)
                draw.rectangle(boxes.tolist(), outline="red")
                draw.rectangle(target_box.tolist(), outline="yellow")
                plt.imshow(img)
                plt.savefig("../result/{}.jpg".format(nnn))
                nnn += 1
            del x, y, category, axes, index, target
            # print(index)
        print("检测到：", total, total_iou / nnn)


if __name__ == "__main__":
    a = Trainer()
    # a.train()
    a.test()