您好,欢迎访问代理记账网站
移动应用 微信公众号 联系我们

咨询热线 -

电话 15988168888

联系客服
  • 价格透明
  • 信息保密
  • 进度掌控
  • 售后无忧

目标识别程序

from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
import os

import torch.nn as nn

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 = 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(".")]取出该值:
        # print([:.index(".")])
        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[: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))
        # 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

"""


class MyNetWork(nn.Module):
    def __init__(self):
        super(MyNetWork, self).__init__()
        self.convlution_layer = nn.Sequential(

            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


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):
        transforms = trans.Compose([
            trans.ToTensor(),
            trans.Normalize(MyDataset.mean,MyDataset.std)
        ])
        self.train_dataset = MyDataset(root=r"../test_train", train=True, transforms=transforms)
        self.test_dataset = MyDataset(root=r"../test_picture", train=False, transforms=transforms)
        self.net = MyNetWork().cuda()
        # 回归问题最常用的损失函数是均方误差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(), lr=1e-5)

    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=64, 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()
                axes, category = self.net(x)  # 前向传播输出坐标和标签
                # (106,)   (106,)
                loss1 = self.offset_lossfunc(axes, y[:,0:4])
                loss2 = self.category_lossfunc(category, y[:,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()

        torch.save(self.net, "../models/net22.pth")


    def test(self):
        testloader = DataLoader(dataset=self.test_dataset, batch_size=100, shuffle=False)
        self.net = torch.load("models/net22.pth")
        total = 0
        nnn = 0
        total_iou = 0
        ii=0
        for x, y in testloader:
            # print(x.size(),y.size())
            ii+=1
            print(ii)
            x, y = x.cuda(), y.cuda()
            # print(y)
            axes, category = self.net(x)
            total += (category.round() == y[:, 4]).sum()  # 预测值等于标签的总数
            # print(y[:, 4])
            index = category.round() == 1
            # print("index",index)
            # print(category)
            """
            这里表示有小黄人的图片的索引集 (最后结果是True 和 False的集合)
            形如:tensor([True, True, False, True, False, True, True, False, True, True])
            """
            target = y[index]  # 有小黄人的图片的标签(包括坐标和分类标签)
            # print(y[0],target[0])
            # print(target.size(),target)
            # print("target:",target[0,0:4])
            """
            还原有小黄人的图片,因为现在要可视化图片,所以要把之前对图片进行的归一化和去均值操作逆向还原回去。
            数据预处理的时候对其做了标准化:处理后的图片=(原始img/255 - mean)/ std  那么现在计算原始图片,原始img =(处理后的图片 * std  + mean)*255
            """
            x = (x[index].cpu() * MyDataset.std.reshape(-1, 3, 1, 1) + MyDataset.mean.reshape(-1, 3, 1, 1))
            # 还原预测为正样本的数据。不用乘以255.。trans.ToPILImage("RGB"):自动会乘以255

            for j, i in enumerate(axes[index]):  # j 为enumerate自动产生的索引
                # print(j,i)
                boxes = (i.data.cpu().numpy() * 224).astype(np.int32)  # 还原预测坐标并将其转化为无符号整型
                # print(boxes)
                target_box = (target[j, 0:4].data.cpu().numpy() * 224).astype(np.int32)  # 还原目标坐标并将其转化为无符号整型
                # print(target_box)
                img = trans.ToPILImage()(x[j])  # 转换图片

                total_iou+=countIOU(boxes,target_box)

                """
                torchvision.transforms.ToPILImage
                对于一个Tensor的转化过程是:
                1. 将张量的每个元素乘上255
                2. 将张量的数据类型有FloatTensor转化成Uint8
                3. 将张量转化成numpy的ndarray类型
                4. 对ndarray对象做transpose (1, 2, 0)的操作
                5. 利用Image下的fromarray函数,将ndarray对象转化成PILImage形式
                6. 输出PILImage
                """
                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
                # plt.pause(0.1)

                # de = rectangle_to_polygon(boxes.tolist())
                # gt = rectangle_to_polygon(target_box.tolist())
                # iou = get_intersection_over_union(de, gt)
                # total_iou += iou

            del x, y, category, axes, index, target
        print("正确率:", total,total_iou/nnn)

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


分享:

低价透明

统一报价,无隐形消费

金牌服务

一对一专属顾问7*24小时金牌服务

信息保密

个人信息安全有保障

售后无忧

服务出问题客服经理全程跟进