2016年提出的Faster RCNN目标检测模型是深度学习现代目标检测算法的开山之作，也是第一个真正全流程都是神经网络的目标检测模型。
其主要步骤如下：
1，使用CNN对输入图片提取feature map.
2，对feature map上的每个点设计一套不同大小和长宽比的anchor作为先验框。
3，设计RPN网络从大量的anchor中筛选出一些作为目标框的proposals并用回归分支纠正它们的位置。
4，使用ROI Pooling技术对不同大小的proposals获取相同大小的对应特征图，以便后续分类模型一并处理。
5，在proposals的feature map上使用分类分支和回归分支进一步预测目标类别和更精确的定位。
anchor技巧和ROI Pooling技术 是非常值得学习的技巧，在许多目标检测模型中都能看到他们的身影。
尽管FasterRCNN历史悠久，但依然是一个非常重要的目标检测任务的baseline.
一般会把它叫做two-stage的目标检测模型，主要是如果train from scratch, RPN网络提取proposals和后续对propasals的定位分类 这两个步骤是要分开训练的，但在微调的时候，通常可以一起训练。
本文我们主要演示调用torchvision中的faster-rcnn模型在自己的数据集上微调来检测螺丝螺母。
#!pip install torchvision,torchkerasimport numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from PIL import Image,ImageColor,ImageDraw,ImageFont
import torch
from torch import nn
import torchvision
from torchvision import datasets, models, transforms
import datetime
import os
import copy
import json
print(torch.__version__)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')2.0.0+cu117〇，预训练模型
from torchkeras.data import get_example_image
img = get_example_image('park.jpg')
img.save('park.jpg')from torchkeras.plots import vis_detection
# 准备数据
inputs = []
img = Image.open('park.jpg').convert("RGB")
img_tensor = torch.from_numpy(np.array(img)/255.).permute(2,0,1).float()
if torch.cuda.is_available():
img_tensor = img_tensor.cuda()
inputs.append(img_tensor)
# 加载模型
num_classes = 91
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
weights=torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights.COCO_V1,
num_classes = num_classes)
if torch.cuda.is_available():
model.to("cuda:0")
model.eval()
# 预测结果
with torch.no_grad():
predictions = model(inputs)
# 结果可视化
class_names = torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights.COCO_V1.meta['categories']
vis_detection(img,predictions[0],class_names,min_score = 0.8)
下面代码我们演示使用我开发的优雅的torchkeras工具在自己的数据集上对Faster-RCNN模型进行finetune。
我们使用一个非常简单的螺丝(bolt)螺母(nut)数据集作为示范。
公众号 算法美食屋 后台回复关键词：torchkeras，获取本文notebook代码和 bolt nut 数据集 下载地址。
一，准备数据
data_path = "./data/bolt_nut"
train_images_path = "./data/bolt_nut/train"
train_targets_path = './data/bolt_nut/train.txt'
val_images_path = "./data/bolt_nut/val"
val_targets_path = './data/bolt_nut/val.txt'
class_names = ['__background__','bolt','nut']class BoltNut(torch.utils.data.Dataset):
def __init__(self, images_path, targets_path,
class_names = class_names,
transforms = None
):
self.images_path = images_path
self.targets_path = targets_path
self.transforms = transforms
self.infos_list = open(targets_path,"r").readlines()
self.class_names = class_names
def __getitem__(self, idx):
info_str = self.infos_list[idx]
info_arr = info_str.replace("\n","").replace("\t ","").split("\t")
img_path = info_arr.pop(0)
info_arr = [x for x in info_arr if x.strip()]
infos = [json.loads(x) for x in info_arr]
img= Image.open(os.path.join(self.images_path,img_path)).convert("RGB")
target = {}
target["image_id"] = torch.tensor([int(img_path.split(".")[0])],dtype = torch.int64)
target["labels"] = torch.tensor([self.class_names.index(x["value"]) for x in infos],
dtype = torch.int64)
coords = [x["coordinate"] for x in infos]
boxes = torch.tensor([[xmin,ymin,xmax,ymax] for (xmin,ymin), (xmax,ymax) in coords])
target["boxes"] = boxes
target["area"] = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
target["iscrowd"] = torch.zeros((len(infos) ,), dtype=torch.int64)
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target
def __len__(self):
return len(self.infos_list)# 可视化数据集
ds_train = BoltNut(train_images_path,train_targets_path)
img,target = ds_train[12]
target["scores"] = torch.ones_like(target["labels"])
img_result = vis_detection(img,target,class_names,min_score = 0.8)
img_result
下面我们设计数据增强模块
import random
from torchvision import transforms as T
class Compose(object):
def __init__(self, transforms):
self.transforms = transforms
def __call__(self, image, target):
for t in self.transforms:
image, target = t(image, target)
return image, target
class RandomHorizontalFlip(object):
def __init__(self, prob):
self.prob = prob
def __call__(self, image, target):
if random.random() < self.prob:
height, width = image.shape[-2:]
image = image.flip(-1)
bbox = target["boxes"]
bbox[:, [0, 2]] = width - bbox[:, [2, 0]]
target["boxes"] = bbox
if "masks" in target:
target["masks"] = target["masks"].flip(-1)
return image, target
class ToTensor(object):
def __call__(self, image, target):
image = T.ToTensor()(image)
return image, targettransforms_train = Compose([ToTensor(),RandomHorizontalFlip(0.5)])
transforms_val = ToTensor()
ds_train = BoltNut(train_images_path,train_targets_path,transforms=transforms_train)
ds_val = BoltNut(val_images_path,val_targets_path,transforms=transforms_val)def collate_fn(batch):
return tuple(zip(*batch))
dl_train = torch.utils.data.DataLoader(ds_train, batch_size=2,
shuffle=True, num_workers=4,collate_fn= collate_fn)
dl_val = torch.utils.data.DataLoader(ds_val, batch_size=2,
shuffle=True, num_workers=4,collate_fn= collate_fn)for batch in dl_train:
features,labels = batch
break二，定义模型
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
num_classes = 3 # 3 classes (bult,nut) + background
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
weights=torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights.COCO_V1)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)三，训练模型
from torchkeras import KerasModelclass StepRunner:
def __init__(self, net, loss_fn, accelerator, stage = "train", metrics_dict = None,
optimizer = None, lr_scheduler = None
):
self.net,self.loss_fn,self.metrics_dict,self.stage = net,loss_fn,metrics_dict,stage
self.optimizer,self.lr_scheduler = optimizer,lr_scheduler
self.accelerator = accelerator
if self.stage=='train':
self.net.train()
else:
self.net.train() #attention here
def __call__(self, batch):
features,labels = batch
#loss
loss_dict = self.net(features,labels)
loss = sum(loss_dict.values())
#backward()
if self.optimizer is not None and self.stage=="train":
self.accelerator.backward(loss)
self.optimizer.step()
if self.lr_scheduler is not None:
self.lr_scheduler.step()
self.optimizer.zero_grad()
#all_preds = self.accelerator.gather(preds)
#all_labels = self.accelerator.gather(labels)
all_loss = self.accelerator.gather(loss).sum()
#losses
step_losses = {self.stage+"_loss":all_loss.item()}
#metrics
step_metrics = {}
if self.stage=="train":
if self.optimizer is not None:
step_metrics['lr'] = self.optimizer.state_dict()['param_groups'][0]['lr']
else:
step_metrics['lr'] = 0.0
return step_losses,step_metrics
KerasModel.StepRunner = StepRunnerparams = [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)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,T_max=4)
keras_model = KerasModel(model,
loss_fn = None,
metrics_dict=None,
optimizer= optimizer,
lr_scheduler=lr_scheduler
)
keras_model.fit(train_data=dl_train,val_data=dl_val,
epochs=20,patience=5,
monitor='val_loss',
mode='min',
ckpt_path ='faster-rcnn.pt',
plot=True
)
四，评估模型
import torch
from PIL import Image
from tqdm import tqdm
from ultralytics.yolo.utils import set_logging
set_logging(verbose=False)
from ultralytics.yolo.utils.metrics import DetMetrics, box_ioudef process_batch(predictions, targets,
iouv = torch.linspace(0.5, 0.95, 10) # iou vector for mAP@0.5:0.95
):
...
return metricsmodel.eval()
list_predictions = [model(x[0].to('cuda')[None,...])[0] for x in ds_val]list_targets = [x[1] for x in ds_val]
names = {0:'bolt',1:'nut'}
metrics = eval_metrics(list_predictions,
list_targets,
names = names)display(metrics.results_dict){'metrics/precision(B)': 0.9976781395819151,
'metrics/recall(B)': 1.0,
'metrics/mAP50(B)': 0.995,
'metrics/mAP50-95(B)': 0.8542317510036526,
'fitness': 0.8683085759032874}import pandas as pd
df = pd.DataFrame()
df['metric'] = metrics.keys
for i,c in names.items():
df[c] = metrics.class_result(i)df
五，使用模型
# 准备数据
inputs = []
img_path = os.path.join(val_images_path,os.listdir(val_images_path)[5])
img = Image.open(img_path).convert("RGB")
img_tensor = torch.from_numpy(np.array(img)/255.).permute(2,0,1).float()
if torch.cuda.is_available():
img_tensor = img_tensor.cuda()
inputs.append(img_tensor)
model.eval()
# 预测结果
with torch.no_grad():
predictions = model(inputs)
# 结果可视化
vis_detection(img,predictions[0],list(idx2names.values()),min_score = 0.8)
公众号 算法美食屋 后台回复关键词：torchkeras，获取本文notebook代码和 bolt nut 数据集 下载地址。
万水千山总是情，点个赞赞行不行？😋😋
