Object Detection using Pytorch
作者:禅与计算机程序设计艺术
1.简介
目标识别是计算机视觉的核心任务之一。该系统能够通过分析图片或视频中的元素对它们进行解析与归类。当前深度学习技术在目标识别领域展现出显著的应用前景。现有的主流目标识别框架多以深度神经网络为基础构建。本文将围绕PyTorch框架展开讨论。
2.基本概念术语说明
- 目标检测 :是指从一副图片或视频中识别出并标记出其中所有出现的特定目标对象,这些对象可能是一个或多个物体、场景中的景物、动物甚至文字等。一般来说,目标检测包含两步:第一步为分类,即对输入图片进行分类,确定输入图片中是否存在感兴趣的目标;第二步为定位,即计算出每个目标对象的具体位置。
- 深度学习 :是一种机器学习方法,它可以从大量的数据中学习到有意义的特征表示,并据此对数据进行预测和分析。深度学习方法主要由深层神经网络构成,它通过前向传播来处理输入数据并输出预测结果。
- 卷积神经网络(CNN) :是一种深度学习网络,它包含卷积层、池化层、归一化层和激活函数层等元素。CNN在图像检测领域占据着举足轻重的地位。
- 锚框(anchor box) :是用于生成候选区域的一种方式。它是在网格边界上采样得到的一组“锚点”,然后将锚点周围一定大小的窗口滑过,得到不同比例和宽高比的窗口作为候选区域。
- 边界框(bounding box) :是指由四个坐标值组成的矩形框,用来表示目标对象的位置、尺寸以及方向角度。
3.核心算法原理及具体操作步骤
3.1 数据集准备
在处理过程中,在实际应用中,数据准备是一个关键步骤。其中所采用的数据集为PASCAL VOC 2007,在计算机视觉领域具有重要地位。该数据集常被用作目标检测领域的基准数据集。它包含20个类别,并为每个类别分配了约200张训练图像和500张测试图像。
import os
from PIL import Image
import xml.etree.ElementTree as ET
class VOCDataset:
def __init__(self, root_dir):
self.root_dir = root_dir
def get_image_annotations(self, image_id):
annotation_file = os.path.join(self.root_dir, "Annotations", "{}.xml".format(image_id))
tree = ET.parse(annotation_file)
root = tree.getroot()
boxes = []
labels = []
for object in root.findall("object"):
class_name = object.find("name").text
if class_name not in ["person"]:
continue
bndbox = object.find("bndbox")
xmin = int(bndbox.find("xmin").text) - 1 # 减1是因为原标注格式的原因
ymin = int(bndbox.find("ymin").text) - 1
xmax = int(bndbox.find("xmax").text)
ymax = int(bndbox.find("ymax").text)
boxes.append([xmin, ymin, xmax, ymax])
label = self.class_to_label[class_name]
labels.append(label)
return np.array(boxes), np.array(labels)
def load_data(self):
data_list = []
images_dir = os.path.join(self.root_dir, "JPEGImages")
annotations_dir = os.path.join(self.root_dir, "Annotations")
self.class_names = sorted(os.listdir(images_dir))
self.class_to_label = {cls_name: i+1 for i, cls_name in enumerate(self.class_names)}
self.num_classes = len(self.class_names)
for filename in os.listdir(images_dir):
basename, _ = os.path.splitext(filename)
img_filepath = os.path.join(images_dir, filename)
anno_filepath = os.path.join(annotations_dir, "{}.xml".format(basename))
try:
with open(anno_filepath, "r", encoding="utf-8") as f:
xmlstr = f.read()
except UnicodeDecodeError:
print("Cannot decode %s" % anno_filepath)
continue
boxes, labels = self.get_image_annotations(basename)
if len(boxes) == 0:
continue
item = {"img_filepath": img_filepath, "boxes": boxes, "labels": labels}
data_list.append(item)
return data_list
dataset = VOCDataset("/home/user/vocdevkit/")
train_data = dataset.load_data()
代码解读
3.2 模型搭建
接下来,我们构建一个基于卷积神经网络(CNN)的目标检测系统。该系统采用YOLO v3算法作为其核心模块,这种算法在目标检测领域具有广泛应用。YOLO v3的网络架构如图所示:该网络主要由二维卷积层、下采样层以及全连接神经元层三大部分构成。
YOLO v3共分为五个部分:
输入模块:接收尺寸为416\times 416像素图像作为输入数据。
基础特征提取模块:通过三层大小均为3\times 3的空间卷积操作生成尺寸缩减至13\times 13的第一级特征图。
置信度预测模块:采用一个单通道(数量为一)的空间卷积单元生成各目标类别的置信度分数(confidence score),该模块共有三个输出通道分别对应物体、人与车辆类别。
边界框回归模块:应用三层空间卷积运算结合两组先验框(anchor boxes)实现对目标边界框位置参数的学习,并生成同样尺寸为13\times 13的第一级特征图。
最终检测模块:通过双通道空间卷积操作将上一层特征图的高度与宽度进一步压缩至图像尺寸的一半即h=7, w=7, 最终获得目标物体定位及其所属类别的概率信息。
使用PyTorch实现YOLO v3的目标检测模型如下所示:
import torch
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
def get_model():
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
num_classes = 3 + 1 # 3 classes (person, car and bike) + background
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
return model
代码解读
3.3 数据加载器
为了加快模型的训练速度, 合理选择数据加载器是提升训练效率的关键. 通常情况下, 当数据集规模较小时, 推荐采用批处理大小为1的小批量随机梯度下降优化算法; 而面对更大规模的数据集时, 建议采用SGD优化方法以加快收敛速度. 对于YOLO v3模型, 其数据加载器的具体实现如下所示:
from PIL import Image
import numpy as np
import cv2
from torch.utils.data import Dataset, DataLoader
class VOCLoader(Dataset):
def __init__(self, data_list, transform=None):
super().__init__()
self.transform = transform
self.data_list = data_list
def __len__(self):
return len(self.data_list)
def __getitem__(self, idx):
item = self.data_list[idx]
img_filepath = item["img_filepath"]
boxes = item["boxes"]
labels = item["labels"].astype(np.int64)
image = cv2.imread(img_filepath)
height, width = image.shape[:2]
original_size = [width, height]
targets = {}
targets['boxes'] = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
targets['labels'] = torch.as_tensor(labels, dtype=torch.int64)
targets['image_id'] = torch.tensor([idx])
targets['area'] = (targets['boxes'][:, 3] - targets['boxes'][:, 1]) * \
(targets['boxes'][:, 2] - targets['boxes'][:, 0])
targets['iscrowd'] = torch.zeros((len(targets['boxes']),), dtype=torch.int64)
if self.transform is not None:
sample = {'image': image, 'bboxes': targets['boxes'],
'labels': targets['labels']}
augmented = self.transform(**sample)
image, targets['boxes'], targets['labels'] = augmented['image'], \
augmented['bboxes'], \
augmented['labels']
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = Image.fromarray(image)
image = transforms.functional.to_tensor(image)
return image, targets, original_size
代码解读
3.4 损失函数及优化器设置
YOLO v3模型由两部分组成:类别置信度损失和边界框定位损失。其中类别置信度损失用于衡量模型预测各类别对象的概率与实际标签之间的差异程度;而边界框定位损失则用于衡量模型预测的边界框位置与真实边界框之间的差异程度。在训练过程中,YOLO v3模型使用SGD(Stochastic Gradient Descent)算法进行优化。具体而言,在训练过程中会通过对前面各层参数进行精细调整实现对前面各层参数的有效微调(fine-tuning),仅更新最后全连接层的所有权重参数。值得注意的是,由于YOLO v3模型具有较高的复杂性特征,在实际应用中其训练过程可能会耗费较多的时间资源。
import torch
import torchvision.transforms as transforms
from utils import VOCLoader, get_model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'device: {device}')
transform = transforms.Compose([
transforms.Resize((416, 416)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
dataset = VOCLoader(train_data, transform=transform)
loader = DataLoader(dataset, batch_size=8, shuffle=True, collate_fn=collate_fn)
model = get_model().to(device)
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
loss_fn = torch.nn.BCEWithLogitsLoss()
mse_loss = torch.nn.MSELoss()
代码解读
3.5 训练过程
在训练阶段中,应定期记录训练日志,其中包括损失函数值.准确率和召回率等关键指标.通过持续追踪这些数据,我们可以有效监控模型在训练数据集上的性能变化趋势,从而判断其是否已达到最优收敛状态.当发现模型在验证数据集上的性能指标达到峰值时,我们应当及时转而对独立的测试数据集进行最终表现评估
from tqdm import tqdm
import time
for epoch in range(10):
total_loss = 0.0
start_time = time.time()
train_loader = loader
model.train()
with tqdm(total=len(loader)) as t:
for iter, (inputs, target, _) in enumerate(train_loader):
inputs = list(image.to(device) for image in inputs)
target = [{k: v.to(device) for k, v in t.items()} for t in target]
optimizer.zero_grad()
outputs = model(inputs)
loss_dict = criterion(outputs, target)
losses = sum(loss for loss in loss_dict.values())
losses.backward()
optimizer.step()
total_loss += float(losses)
t.update()
avg_loss = total_loss / len(train_loader)
end_time = time.time()
print(f'[Epoch {epoch+1}] Average Loss={avg_loss:.4f}, Time Elapsed={(end_time - start_time)/60:.2f} minutes')
代码解读
4. 评估与可视化
在训练完成后,我们可以通过对测试集上预测结果的评估来验证模型的准确性。具体而言,在实现这一目标的过程中,我们可以构建一个混淆矩阵,并计算包括精确度、召回率、F1 Score以及平均IOU等关键指标的具体数值以辅助分析。此外,在完成模型训练后还能够通过数据可视化的方式展示预测结果的分布情况,并对模型的表现情况进行直观的观察和判断
import matplotlib.pyplot as plt
import seaborn as sn
def plot_confusion_matrix(cm, title='Confusion matrix', cmap=plt.cm.Blues):
plt.imshow(cm, interpolation='nearest', cmap=cmap)
plt.title(title)
plt.colorbar()
tick_marks = np.arange(len(class_names))
plt.xticks(tick_marks, class_names, rotation=45)
plt.yticks(tick_marks, class_names)
fmt = '.2f'
thresh = cm.max() / 2.
for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
plt.text(j, i, format(cm[i, j], fmt), horizontalalignment='center', verticalalignment='center',
color='white' if cm[i, j] > thresh else 'black')
plt.tight_layout()
plt.ylabel('True label')
plt.xlabel('Predicted label')
def evaluate(model, test_loader):
y_true = []
y_pred = []
gt_boxes = []
pred_boxes = []
with torch.no_grad():
model.eval()
for inputs, _, original_sizes in test_loader:
input_batch = list(image.to(device) for image in inputs)
output_batch = model(input_batch)
probas = output_batch.softmax(dim=1)
labels = output_batch.labels
for i, (probas_, labels_) in enumerate(zip(probas, labels)):
h, w = original_sizes[i][0], original_sizes[i][1]
for box, label in zip(output_batch[i]['boxes'], output_batch[i]['labels']):
x1 = max(min(round(box[0].item()), w), 0)
y1 = max(min(round(box[1].item()), h), 0)
x2 = min(max(round(box[2].item()), 0), w)
y2 = min(max(round(box[3].item()), 0), h)
if label!= 0:
gt_boxes.append([x1, y1, x2, y2])
y_true.append(label.item()-1)
for score, label in zip(probas_[0:-1], labels_[0:-1]):
if score >= 0.1:
x1 = max(min(round(output_batch[i]['boxes'][label][0].item()*w), w), 0)
y1 = max(min(round(output_batch[i]['boxes'][label][1].item()*h), h), 0)
x2 = min(max(round(output_batch[i]['boxes'][label][2].item()*w), 0), w)
y2 = min(max(round(output_batch[i]['boxes'][label][3].item()*h), 0), h)
pred_boxes.append([x1, y1, x2, y2])
y_pred.append(label.item())
cm = confusion_matrix(y_true, y_pred)
acc = accuracy_score(y_true, y_pred)
recall = recall_score(y_true, y_pred, average='weighted')
precision = precision_score(y_true, y_pred, average='weighted')
f1_score = f1_score(y_true, y_pred, average='weighted')
cf = pd.DataFrame(cm, index=class_names[:-1], columns=class_names[:-1])
fig = plt.figure(figsize=(10, 7))
ax = sn.heatmap(cf, annot=True, fmt='g')
ax.set(xlabel='Predicted Label', ylabel='Ground Truth Label', title=f'{acc*100:.2f}% ({recall:.3f})')
plt.show()
im_resized = cv2.resize(im, (im.shape[1]*3, im.shape[0]*3))
for box in gt_boxes:
cv2.rectangle(im_resized, (box[0]-5, box[1]-5), (box[2]+5, box[3]+5), (255, 0, 0), 1)
for box in pred_boxes:
cv2.rectangle(im_resized, (box[0]-5, box[1]-5), (box[2]+5, box[3]+5), (0, 255, 0), 1)
cv2.imshow('Result', im_resized)
cv2.waitKey()
cv2.destroyAllWindows()
metrics = OrderedDict({'Accuracy': acc,
'Recall': recall,
'Precision': precision,
'F1 Score': f1_score})
return metrics
代码解读
