【PyTorch】单目标检测项目

两种部署情况：部署在 PyTorch 数据集上，以及部署在本地存储的单个映像上。

目录

定义数据集

搭建模型

部署模型

---

定义数据集

详细参照前文【PyTorch】单目标检测项目

import torchvision
import os
import pandas as pd
import matplotlib.pylab as plt
import torch
%matplotlib inline# 定义一个函数，用于将a列表中的元素除以b列表中的对应元素，返回一个新的列表
def scale_label(a,b):# 使用zip函数将a和b列表中的元素一一对应div = [ai/bi for ai,bi in zip(a,b)]# 返回新的列表return div# 定义一个函数，用于将a列表中的元素乘以b列表中的对应元素，返回一个新的列表
def rescale_label(a,b):# 使用zip函数将a和b列表中的元素一一对应div = [ai*bi for ai,bi in zip(a,b)]# 返回新的列表return divimport torchvision.transforms.functional as TF# 定义一个函数，用于调整图像和标签的大小
def resize_img_label(image,label=(0.,0.),target_size=(256,256)):# 获取原始图像的宽度和高度w_orig,h_orig = image.size   # 获取目标图像的宽度和高度w_target,h_target = target_size# 获取标签的坐标cx, cy= label# 调整图像大小image_new = TF.resize(image,target_size)# 调整标签大小label_new= cx/w_orig*w_target, cy/h_orig*h_target# 定义一个transformer函数，用于对图像和标签进行变换
def transformer(image, label, params):# 调用resize_img_label函数，对图像和标签进行尺寸调整image,label=resize_img_label(image,label,params["target_size"])# 如果params中scale_label参数为True，则调用scale_label函数，对标签进行缩放if params["scale_label"]:label=scale_label(label,params["target_size"])# 将图像转换为张量image=TF.to_tensor(image)# 返回变换后的图像和标签return image, labelfrom torch.utils.data import Dataset
from PIL import Imageclass AMD_dataset(Dataset):def __init__(self, path2data, transform, trans_params):      # 初始化函数，传入数据路径、转换函数和转换参数pass    def __len__(self):# 返回数据集的大小return len(self.labels)def __getitem__(self, idx):# 根据索引获取数据集中的一个样本pass# 返回调整后的图像和标签return image_new,label_newdef __init__(self, path2data, transform, trans_params):      # 获取标签文件路径path2labels=os.path.join(path2data,"Training400","Fovea_location.xlsx")# 读取标签文件labels_df=pd.read_excel(path2labels,engine='openpyxl',index_col="ID")# 获取标签数据self.labels = labels_df[["Fovea_X","Fovea_Y"]].values# 获取图片名称self.imgName=labels_df["imgName"]# 获取图片IDself.ids=labels_df.index# 获取图片全路径self.fullPath2img=[0]*len(self.ids)for id_ in self.ids:# 根据图片名称判断图片类型if self.imgName[id_][0]=="A":prefix="AMD"else:prefix="Non-AMD"# 获取图片全路径self.fullPath2img[id_-1]=os.path.join(path2data,"Training400",prefix,self.imgName[id_])# 获取数据转换函数self.transform = transform# 获取数据转换参数self.trans_params=trans_paramsdef __getitem__(self, idx):# 打开指定索引的图像image = Image.open(self.fullPath2img[idx])  # 获取指定索引的标签label= self.labels[idx]# 对图像和标签进行变换image,label = self.transform(image,label,self.trans_params)# 返回变换后的图像和标签return image, label#重写
AMD_dataset.__init__=__init__
AMD_dataset.__getitem__=__getitem__path2data="./data/"#验证参数
trans_params_val={"target_size" : (256, 256),"p_hflip" : 0.0,"p_vflip" : 0.0,"p_shift" : 0.0,"p_brightness": 0.0,"p_contrast": 0.0,"p_gamma": 0.0,"gamma": 0.0,"scale_label": True,    
}
amd_ds2=AMD_dataset(path2data,transformer,trans_params_val)from sklearn.model_selection import ShuffleSplitsss = ShuffleSplit(n_splits=1, test_size=0.2, random_state=0)indices=range(len(amd_ds2))for train_index, val_index in sss.split(indices):print(len(train_index))print("-"*10)print(len(val_index))from torch.utils.data import Subset# 创建一个Subset对象，将amd_ds2数据集按照val_index索引进行划分，得到验证集val_ds
val_ds=Subset(amd_ds2,val_index)from torch.utils.data import DataLoader
# 创建一个DataLoader对象，用于加载验证集数据
val_dl = DataLoader(val_ds, batch_size=16, shuffle=False)  

搭建模型

 详细参照前文【PyTorch】单目标检测项目

import torch.nn as nn
import torch.nn.functional as Fclass Net(nn.Module):def __init__(self, params):super(Net, self).__init__()def forward(self, x):return xdef __init__(self, params):super(Net, self).__init__()C_in,H_in,W_in=params["input_shape"]init_f=params["initial_filters"] num_outputs=params["num_outputs"] self.conv1 = nn.Conv2d(C_in, init_f, kernel_size=3,stride=2,padding=1)self.conv2 = nn.Conv2d(init_f+C_in, 2*init_f, kernel_size=3,stride=1,padding=1)self.conv3 = nn.Conv2d(3*init_f+C_in, 4*init_f, kernel_size=3,padding=1)self.conv4 = nn.Conv2d(7*init_f+C_in, 8*init_f, kernel_size=3,padding=1)self.conv5 = nn.Conv2d(15*init_f+C_in, 16*init_f, kernel_size=3,padding=1)self.fc1 = nn.Linear(16*init_f, num_outputs)def forward(self, x):identity=F.avg_pool2d(x,4,4)x = F.relu(self.conv1(x))x = F.max_pool2d(x, 2, 2)x = torch.cat((x, identity), dim=1)identity=F.avg_pool2d(x,2,2)x = F.relu(self.conv2(x))x = F.max_pool2d(x, 2, 2)x = torch.cat((x, identity), dim=1)identity=F.avg_pool2d(x,2,2)x = F.relu(self.conv3(x))x = F.max_pool2d(x, 2, 2)x = torch.cat((x, identity), dim=1)identity=F.avg_pool2d(x,2,2)x = F.relu(self.conv4(x))x = F.max_pool2d(x, 2, 2)x = torch.cat((x, identity), dim=1)x = F.relu(self.conv5(x))x=F.adaptive_avg_pool2d(x,1)x = x.reshape(x.size(0), -1)x = self.fc1(x)return xNet.__init__=__init__
Net.forward=forwardparams_model={"input_shape": (3,256,256),"initial_filters": 16, "num_outputs": 2,}
model = Net(params_model)
model.eval()if torch.cuda.is_available():device = torch.device("cuda")model=model.to(device)  

部署模型

path2weights="./models/weights.pt"
# 加载模型权重
model.load_state_dict(torch.load(path2weights))# 定义设备为GPU
device = torch.device("cuda")# 定义一个函数，用于计算模型在数据集上的损失和指标
def loss_epoch(model,loss_func,dataset_dl,sanity_check=False,opt=None):# 初始化运行损失和运行指标running_loss=0.0running_metric=0.0# 获取数据集的长度len_data=len(dataset_dl.dataset)# 遍历数据集for xb, yb in dataset_dl:# 将标签堆叠成一维张量yb=torch.stack(yb,1)# 将标签转换为浮点型，并移动到GPU上yb=yb.type(torch.float32).to(device)# 将输入数据移动到GPU上，并获取模型输出output=model(xb.to(device))# 计算当前批次的损失和指标loss_b,metric_b=loss_batch(loss_func, output, yb, opt)# 累加损失running_loss+=loss_b# 如果指标不为空，则累加指标if metric_b is not None:running_metric+=metric_b# 如果是进行sanity check，则只计算一个批次if sanity_check is True:break# 计算平均损失loss=running_loss/float(len_data)# 计算平均指标metric=running_metric/float(len_data)# 返回平均损失和平均指标return loss, metric# 将中心点坐标和宽高转换为边界框坐标
def cxcy2bbox(cxcy,w=50./256,h=50./256):# 创建一个与cxcy形状相同的张量，每个元素都为ww_tensor=torch.ones(cxcy.shape[0],1,device=cxcy.device)*w# 创建一个与cxcy形状相同的张量，每个元素都为hh_tensor=torch.ones(cxcy.shape[0],1,device=cxcy.device)*h# 将cxcy的第一列提取出来，并增加一个维度cx=cxcy[:,0].unsqueeze(1)# 将cxcy的第二列提取出来，并增加一个维度cy=cxcy[:,1].unsqueeze(1)# 将cx、cy、w_tensor、h_tensor按列拼接起来boxes=torch.cat((cx,cy, w_tensor, h_tensor), -1)  # 将boxes的第一列和第二列分别减去w_tensor和h_tensor的一半，然后将结果按行拼接起来return torch.cat((boxes[:, :2] - boxes[:, 2:]/2,boxes[:, :2] + boxes[:, 2:]/2), 1)  # 定义一个函数metrics_batch，用于计算输出和目标之间的交并比
def metrics_batch(output, target):# 将输出和目标转换为边界框格式output=cxcy2bbox(output)target=cxcy2bbox(target)# 计算输出和目标之间的交并比iou=torchvision.ops.box_iou(output, target)# 返回交并比的和return torch.diagonal(iou, 0).sum().item()# 定义一个函数，用于计算损失函数、输出和目标之间的损失值
def loss_batch(loss_func, output, target, opt=None):# 计算输出和目标之间的损失值loss = loss_func(output, target)# 计算输出和目标之间的度量值metric_b = metrics_batch(output,target)# 如果opt不为空，则执行反向传播和优化if opt is not None:opt.zero_grad()loss.backward()opt.step()# 返回损失值和度量值return loss.item(), metric_b# 定义损失函数，使用SmoothL1Loss，reduction参数设置为sum
loss_func=nn.SmoothL1Loss(reduction="sum")# 在不计算梯度的情况下，计算模型在验证集上的损失和指标
with torch.no_grad():loss,metric=loss_epoch(model,loss_func,val_dl)# 打印损失和指标
print(loss,metric) 

from PIL import ImageDraw
import numpy as np
import torchvision.transforms.functional as tv_F
np.random.seed(0)import matplotlib.pylab as plt
%matplotlib inlinedef show_tensor_2labels(img,label1,label2,w_h=(50,50)): # 将label1和label2按照img的shape进行缩放label1=rescale_label(label1,img.shape[1:])label2=rescale_label(label2,img.shape[1:])# 将img转换为PIL图像img=tv_F.to_pil_image(img) # 获取w_h的宽度和高度w,h=w_h # 获取label1的坐标cx,cy=label1# 在img上绘制一个绿色的矩形draw = ImageDraw.Draw(img)draw.rectangle(((cx-w/2, cy-h/2), (cx+w/2, cy+h/2)),outline="green",width=2)# 获取label2的坐标cx,cy=label2# 在img上绘制一个红色的矩形draw.rectangle(((cx-w/2, cy-h/2), (cx+w/2, cy+h/2)),outline="red",width=2)# 显示imgplt.imshow(np.asarray(img))# 生成一个长度为10的随机整数数组，数组中的元素为0到len(val_ds)之间的随机整数
rndInds=np.random.randint(len(val_ds),size=10)
# 打印生成的随机整数数组
print(rndInds)

# 设置图像大小
plt.rcParams['figure.figsize'] = (15, 10)
# 调整子图之间的间距
plt.subplots_adjust(wspace=0.0, hspace=0.15)# 遍历随机索引
for i,rndi in enumerate(rndInds):# 获取图像和标签img,label=val_ds[rndi]# 获取图像的宽度和高度h,w=img.shape[1:]# 不计算梯度with torch.no_grad():# 获取模型预测的标签label_pred=model(img.unsqueeze(0).to(device))[0].cpu()# 绘制子图plt.subplot(2,3,i+1)# 显示图像和标签show_tensor_2labels(img,label,label_pred)# 将标签转换为边界框label_bb=cxcy2bbox(torch.tensor(label).unsqueeze(0))# 将模型预测的标签转换为边界框label_pred_bb=cxcy2bbox(label_pred.unsqueeze(0))# 计算IOUiou=torchvision.ops.box_iou(label_bb, label_pred_bb)        # 设置标题plt.title("%.2f" %iou.item())# 如果索引大于4，则跳出循环if i>4:break

# 定义一个函数，用于加载图片和标签
def load_img_label(labels_df,id_):    # 获取图片名称imgName=labels_df["imgName"]    # 判断图片名称是否以"A"开头if imgName[id_][0]=="A":# 如果是，则前缀为"AMD"prefix="AMD"else:# 否则，前缀为"Non-AMD"prefix="Non-AMD"# 拼接图片路径fullPath2img=os.path.join(path2data,"Training400",prefix,imgName[id_])# 打开图片img = Image.open(fullPath2img)# 获取图片中心点坐标x=labels_df["Fovea_X"][id_]y=labels_df["Fovea_Y"][id_]# 返回图片和中心点坐标label=(x,y)return img,label# 定义标签文件路径
path2labels=os.path.join(path2data,"Training400","Fovea_location.xlsx")
# 读取标签文件，使用openpyxl引擎，将ID列作为索引
labels_df=pd.read_excel(path2labels,engine='openpyxl',index_col="ID")# 加载图片和标签，使用标签文件中的第一行数据
img,label=load_img_label(labels_df,1)   
# 打印图片和标签的大小
print(img.size, label)# 调整图片和标签的大小为256x256
img,label=resize_img_label(img,label,target_size=(256,256))
# 打印调整后的图片和标签的大小
print(img.size, label)# 将图片转换为张量
img=TF.to_tensor(img)
# 将标签缩放到256x256
label=scale_label(label,(256,256))
# 打印转换后的图片的形状
print(img.shape)# 在不计算梯度的情况下，使用模型对图片进行预测
with torch.no_grad():label_pred=model(img.unsqueeze(0).to(device))[0].cpu()# 显示图片和标签以及预测结果
show_tensor_2labels(img,label,label_pred)

import time
# 定义一个空列表，用于存储每次推理的时间
elapsed_times=[]
# 不计算梯度，进行推理
with torch.no_grad():# 循环100次for k in range(100):# 记录开始时间start=time.time()# 对输入图片进行推理，并获取预测结果label_pred=model(img.unsqueeze(0).to(device))[0].cpu()# 计算推理时间elapsed=time.time()-start# 将每次推理的时间添加到列表中elapsed_times.append(elapsed)
# 打印每次推理的平均时间
print("inference time per image: %.4f s" %np.mean(elapsed_times))