文章目录
前言
本文主要讲解基于mxnet深度学习框架实现目标检测，实现的模型为FasterRCNN
环境配置：
python 3.8
mxnet 1.7.0
cuda 10.1
目标检测发展史及意义
图像分类任务的实现可以让我们粗略的知道图像中包含了什么类型的物体，但并不知道物体在图像中哪一个位置，也不知道物体的具体信息，在一些具体的应用场景比如车牌识别、交通违章检测、人脸识别、运动捕捉，单纯的图像分类就不能完全满足我们的需求了。
这时候，需要引入图像领域另一个重要任务：物体的检测与识别。在传统机器领域，一个典型的案例是利用HOG（Histogram of Gradient）特征来生成各种物体相应的“滤波器”，HOG滤波器能完整的记录物体的边缘和轮廓信息，利用这一滤波器过滤不同图片的不同位置，当输出响应值幅度超过一定阈值，就认为滤波器和图片中的物体匹配程度较高，从而完成了物体的检测。
一、数据集的准备
首先我是用的是halcon数据集里边的药片，去了前边的100张做标注，后面的300张做测试，其中100张里边选择90张做训练集，10张做验证集。
1.标注工具的安装
pip install labelimg
进入cmd，输入labelimg，会出现如图的标注工具：
2.数据集的准备
首先我们先创建3个文件夹，如图：
DataImage：100张需要标注的图像
DataLabel：空文件夹，主要是存放标注文件，这个在labelimg中生成标注文件
test：存放剩下的300张图片，不需要标注
DataImage目录下和test目录的存放样子是这样的（以DataImage为例）：
3.标注数据
首先我们需要在labelimg中设置图像路径和标签存放路径，如图：
然后先记住快捷键：w：开始编辑，a:上一张，d:下一张。这个工具只需要这三个快捷键即可完成工作。
开始标注工作，首先按下键盘w，这个时候进入编辑框框的模式，然后在图像上绘制框框，输入标签（框框属于什么类别），即可完成物体1的标注，一张物体可以多个标注和多个类别，但是切记不可摸棱两可，比如这张图像对于某物体标注了，另一张图像如果出现同样的就需要标注，或者标签类别不可多个，比如这个图象A物体标注为A标签，下张图的A物体标出成了B标签，最终的效果如图：
最后标注完成会在DataLabel中看到标注文件，json格式：
4.解释xml文件的内容
xml标签文件如图，我们用到的就只有object对象，对其进行解析即可。
二、网络结构的介绍
论文地址：https://arxiv.org/pdf/1506.01497.pdf
网络结构：
Faster RCNN其实可以分为4个主要内容：
- Conv layers。作为一种CNN网络目标检测方法，Faster RCNN首先使用一组基础的conv+relu+pooling层提取image的feature maps。该feature maps被共享用于后续RPN层和全连接层。
- Region Proposal Networks。RPN网络用于生成region proposals。该层通过softmax判断anchors属于positive或者negative，再利用bounding box regression修正anchors获得精确的proposals。
- Roi Pooling。该层收集输入的feature maps和proposals，综合这些信息后提取proposal feature maps，送入后续全连接层判定目标类别。
- Classification。利用proposal feature maps计算proposal的类别，同时再次bounding box regression获得检测框最终的精确位置。
三、代码实现
0.工程目录结构如下
core：损失计算及一些核心计算的文件都存放在此文件夹
data：数据加载的相关函数及类
net：包含主干网络结构及标准的fasterrcnn结构
utils：数据预处理的相关文件
Ctu_FasterRCNN.py：fasterrcnn的训练类和测试类，是整个AI的主入口
1.导入库
import os, time, sys, json, colorsys, cv2, copy
sys.path.append('.')
from mxnet import nd
import numpy as np
import mxnet as mx
from PIL import Image, ImageDraw, ImageFont
from mxnet import gluon
from mxnet.contrib import amp
from data.data_loader import VOCDetection,VOC07MApMetric,MixupDetection
from nets.faster_rcnn.predefined_models import faster_rcnn_resnet18_v1b_voc,faster_rcnn_resnet50_v1b_voc, faster_rcnn_resnet101_v1d_voc
from data.batchify_fn import Tuple,Append,FasterRCNNTrainBatchify
from data.sampler import SplitSortedBucketSampler
from data.data_transform import FasterRCNNDefaultTrainTransform, FasterRCNNDefaultValTransform
from core.loss import RPNAccMetric, RPNL1LossMetric, RCNNAccMetric, RCNNL1LossMetric
from nets.faster_rcnn.data_parallel import ForwardBackwardTask
from utils.parallel import Parallel
from nets.nn.bbox import BBoxClipToImage
2.配置GPU/CPU环境
self.ctx = [mx.gpu(int(i)) for i in USEGPU.split(',') if i.strip()]
self.ctx = self.ctx if self.ctx else [mx.cpu()]
3.数据加载器
这里输入的是迭代器，后面都会利用它构建训练的迭代器
class VOCDetection(dataset.Dataset):
def CreateDataList(self,IMGDir,XMLDir):
ImgList = os.listdir(IMGDir)
XmlList = os.listdir(XMLDir)
classes = []
dataList=[]
for each_jpg in ImgList:
each_xml = each_jpg.split('.')[0] + '.xml'
if each_xml in XmlList:
dataList.append([os.path.join(IMGDir,each_jpg),os.path.join(XMLDir,each_xml)])
with open(os.path.join(XMLDir,each_xml), "r", encoding="utf-8") as in_file:
tree = ET.parse(in_file)
root = tree.getroot()
for obj in root.iter('object'):
cls = obj.find('name').text
if cls not in classes:
classes.append(cls)
return dataList,classes
def __init__(self, ImageDir, XMLDir,transform=None):
self.datalist,self.classes_names = self.CreateDataList(ImageDir,XMLDir)
self._transform = transform
self.index_map = dict(zip(self.classes_names, range(len(self.classes_names))))
# self._label_cache = self._preload_labels()
@property
def classes(self):
return self.classes_names
def __len__(self):
return len(self.datalist)
def __getitem__(self, idx):
img_path = self.datalist[idx][0]
# label = self._label_cache[idx] if self._label_cache else self._load_label(idx)
label = self._load_label(idx)
img = mx.image.imread(img_path, 1)
if self._transform is not None:
return self._transform(img, label)
return img, label.copy()
def _preload_labels(self):
return [self._load_label(idx) for idx in range(len(self))]
def _load_label(self, idx):
anno_path = self.datalist[idx][1]
root = ET.parse(anno_path).getroot()
size = root.find('size')
width = float(size.find('width').text)
height = float(size.find('height').text)
label = []
for obj in root.iter('object'):
try:
difficult = int(obj.find('difficult').text)
except ValueError:
difficult = 0
cls_name = obj.find('name').text.strip().lower()
if cls_name not in self.classes:
continue
cls_id = self.index_map[cls_name]
xml_box = obj.find('bndbox')
xmin = (float(xml_box.find('xmin').text) - 1)
ymin = (float(xml_box.find('ymin').text) - 1)
xmax = (float(xml_box.find('xmax').text) - 1)
ymax = (float(xml_box.find('ymax').text) - 1)
try:
self._validate_label(xmin, ymin, xmax, ymax, width, height)
label.append([xmin, ymin, xmax, ymax, cls_id, difficult])
except AssertionError as e:
pass
return np.array(label)
def _validate_label(self, xmin, ymin, xmax, ymax, width, height):
assert 0 <= xmin < width, "xmin must in [0, {}), given {}".format(width, xmin)
assert 0 <= ymin < height, "ymin must in [0, {}), given {}".format(height, ymin)
assert xmin < xmax <= width, "xmax must in (xmin, {}], given {}".format(width, xmax)
assert ymin < ymax <= height, "ymax must in (ymin, {}], given {}".format(height, ymax)
4.模型构建
本项目使用resne做主干网络结构
def faster_rcnn_resnet50_v1b_voc(classes,small_size=600,max_size = 1000, **kwargs):
base_network = resnet50_v1b(classes, dilated=False, use_global_stats=True, **kwargs)
features = nn.HybridSequential()
top_features = nn.HybridSequential()
for layer in ['conv1', 'bn1', 'relu', 'maxpool', 'layer1', 'layer2', 'layer3']:
features.add(getattr(base_network, layer))
for layer in ['layer4']:
top_features.add(getattr(base_network, layer))
train_patterns = '|'.join(['.*dense', '.*rpn', '.*down(2|3|4)_conv', '.*layers(2|3|4)_conv'])
return get_faster_rcnn(
features=features, top_features=top_features, classes=classes,
short=small_size, max_size=max_size, train_patterns=train_patterns,
nms_thresh=0.3, nms_topk=400, post_nms=100,
roi_mode='align', roi_size=(14, 14), strides=16, clip=None,
rpn_channel=1024, base_size=16, scales=(2, 4, 8, 16, 32),
ratios=(0.5, 1, 2), alloc_size=(128, 128), rpn_nms_thresh=0.7,
rpn_train_pre_nms=12000, rpn_train_post_nms=2000,
rpn_test_pre_nms=6000, rpn_test_post_nms=300, rpn_min_size=16,
num_sample=128, pos_iou_thresh=0.5, pos_ratio=0.25, max_num_gt=100,
**kwargs)```
```python
class FasterRCNN(RCNN):
def __init__(self, features, top_features, classes, box_features=None,
short=600, max_size=1000, min_stage=4, max_stage=4, train_patterns=None,
nms_thresh=0.3, nms_topk=400, post_nms=100, roi_mode='align', roi_size=(14, 14), strides=16,
clip=None, rpn_channel=1024, base_size=16, scales=(8, 16, 32),
ratios=(0.5, 1, 2), alloc_size=(128, 128), rpn_nms_thresh=0.7,
rpn_train_pre_nms=12000, rpn_train_post_nms=2000, rpn_test_pre_nms=6000,
rpn_test_post_nms=300, rpn_min_size=16, per_device_batch_size=1, num_sample=128,
pos_iou_thresh=0.5, pos_ratio=0.25, max_num_gt=300, additional_output=False,
force_nms=False, minimal_opset=False, **kwargs):
super(FasterRCNN, self).__init__(
features=features, top_features=top_features, classes=classes,
box_features=box_features, short=short, max_size=max_size,
train_patterns=train_patterns, nms_thresh=nms_thresh, nms_topk=nms_topk, post_nms=post_nms,
roi_mode=roi_mode, roi_size=roi_size, strides=strides, clip=clip, force_nms=force_nms,
minimal_opset=minimal_opset, **kwargs)
if max_stage - min_stage > 1 and isinstance(strides, (int, float)):
raise ValueError('Multi level detected but strides is of a single number:', strides)
if rpn_train_post_nms > rpn_train_pre_nms:
rpn_train_post_nms = rpn_train_pre_nms
if rpn_test_post_nms > rpn_test_pre_nms:
rpn_test_post_nms = rpn_test_pre_nms
self.ashape = alloc_size[0]
self._min_stage = min_stage
self._max_stage = max_stage
self.num_stages = max_stage - min_stage + 1
if self.num_stages > 1:
assert len(scales) == len(strides) == self.num_stages, "The num_stages (%d) must match number of scales (%d) and strides (%d)" % (self.num_stages, len(scales), len(strides))
self._batch_size = per_device_batch_size
self._num_sample = num_sample
self._rpn_test_post_nms = rpn_test_post_nms
if minimal_opset:
self._target_generator = None
else:
self._target_generator = lambda: RCNNTargetGenerator(self.num_class, int(num_sample * pos_ratio), self._batch_size)
self._additional_output = additional_output
with self.name_scope():
self.rpn = RPN(
channels=rpn_channel, strides=strides, base_size=base_size,
scales=scales, ratios=ratios, alloc_size=alloc_size,
clip=clip, nms_thresh=rpn_nms_thresh, train_pre_nms=rpn_train_pre_nms,
train_post_nms=rpn_train_post_nms, test_pre_nms=rpn_test_pre_nms,
test_post_nms=rpn_test_post_nms, min_size=rpn_min_size,
multi_level=self.num_stages > 1, per_level_nms=False,
minimal_opset=minimal_opset)
self.sampler = RCNNTargetSampler(num_image=self._batch_size, num_proposal=rpn_train_post_nms, num_sample=num_sample, pos_iou_thresh=pos_iou_thresh, pos_ratio=pos_ratio, max_num_gt=max_num_gt)
@property
def target_generator(self):
if self._target_generator is None:
raise ValueError("`minimal_opset` enabled, target generator is not available")
if not isinstance(self._target_generator, mx.gluon.Block):
self._target_generator = self._target_generator()
self._target_generator.initialize()
return self._target_generator
def reset_class(self, classes, reuse_weights=None):
super(FasterRCNN, self).reset_class(classes, reuse_weights)
self._target_generator = lambda: RCNNTargetGenerator(self.num_class, self.sampler._max_pos, self._batch_size)
def _pyramid_roi_feats(self, F, features, rpn_rois, roi_size, strides, roi_mode='align', roi_canonical_scale=224.0, sampling_ratio=2, eps=1e-6):
max_stage = self._max_stage
if self._max_stage > 5:
max_stage = self._max_stage - 1
_, x1, y1, x2, y2 = F.split(rpn_rois, axis=-1, num_outputs=5)
h = y2 - y1 + 1
w = x2 - x1 + 1
roi_level = F.floor(4 + F.log2(F.sqrt(w * h) / roi_canonical_scale + eps))
roi_level = F.squeeze(F.clip(roi_level, self._min_stage, max_stage))
pooled_roi_feats = []
for i, l in enumerate(range(self._min_stage, max_stage + 1)):
if roi_mode == 'pool':
pooled_feature = F.ROIPooling(features[i], rpn_rois, roi_size, 1. / strides[i])
pooled_feature = F.where(roi_level == l, pooled_feature, F.zeros_like(pooled_feature))
elif roi_mode == 'align':
if 'box_encode' in F.contrib.__dict__ and 'box_decode' in F.contrib.__dict__:
masked_rpn_rois = F.where(roi_level == l, rpn_rois, F.ones_like(rpn_rois) * -1.)
pooled_feature = F.contrib.ROIAlign(features[i], masked_rpn_rois, roi_size, 1. / strides[i], sample_ratio=sampling_ratio)
else:
pooled_feature = F.contrib.ROIAlign(features[i], rpn_rois, roi_size, 1. / strides[i], sample_ratio=sampling_ratio)
pooled_feature = F.where(roi_level == l, pooled_feature, F.zeros_like(pooled_feature))
else:
raise ValueError("Invalid roi mode: {}".format(roi_mode))
pooled_roi_feats.append(pooled_feature)
pooled_roi_feats = F.ElementWiseSum(*pooled_roi_feats)
return pooled_roi_feats
def hybrid_forward(self, F, x, gt_box=None, gt_label=None):
def _split(x, axis, num_outputs, squeeze_axis):
x = F.split(x, axis=axis, num_outputs=num_outputs, squeeze_axis=squeeze_axis)
if isinstance(x, list):
return x
else:
return [x]
feat = self.features(x)
if not isinstance(feat, (list, tuple)):
feat = [feat]
if autograd.is_training():
rpn_score, rpn_box, raw_rpn_score, raw_rpn_box, anchors = self.rpn(F.zeros_like(x), *feat)
rpn_box, samples, matches = self.sampler(rpn_box, rpn_score, gt_box)
else:
_, rpn_box = self.rpn(F.zeros_like(x), *feat)
num_roi = self._num_sample if autograd.is_training() else self._rpn_test_post_nms
batch_size = self._batch_size if autograd.is_training() else 1
with autograd.pause():
roi_batchid = F.arange(0, batch_size)
roi_batchid = F.repeat(roi_batchid, num_roi)
rpn_roi = F.concat(*[roi_batchid.reshape((-1, 1)), rpn_box.reshape((-1, 4))], dim=-1)
rpn_roi = F.stop_gradient(rpn_roi)
if self.num_stages > 1:
pooled_feat = self._pyramid_roi_feats(F, feat, rpn_roi, self._roi_size, self._strides, roi_mode=self._roi_mode)
else:
if self._roi_mode == 'pool':
pooled_feat = F.ROIPooling(feat[0], rpn_roi, self._roi_size, 1. / self._strides)
elif self._roi_mode == 'align':
pooled_feat = F.contrib.ROIAlign(feat[0], rpn_roi, self._roi_size, 1. / self._strides, sample_ratio=2)
else:
raise ValueError("Invalid roi mode: {}".format(self._roi_mode))
if self.top_features is not None:
top_feat = self.top_features(pooled_feat)
else:
top_feat = pooled_feat
if self.box_features is None:
box_feat = F.contrib.AdaptiveAvgPooling2D(top_feat, output_size=1)
else:
box_feat = self.box_features(top_feat)
cls_pred = self.class_predictor(box_feat)
cls_pred = cls_pred.reshape((batch_size, num_roi, self.num_class + 1))
if autograd.is_training():
cls_targets, box_targets, box_masks, indices = self.target_generator(rpn_box, samples, matches, gt_label, gt_box)
box_feat = F.reshape(box_feat.expand_dims(0), (batch_size, -1, 0))
box_pred = self.box_predictor(F.concat(
*[F.take(F.slice_axis(box_feat, axis=0, begin=i, end=i + 1).squeeze(), F.slice_axis(indices, axis=0, begin=i, end=i + 1).squeeze())
for i in range(batch_size)], dim=0))
box_pred = box_pred.reshape((batch_size, -1, self.num_class, 4))
if self._additional_output:
return (cls_pred, box_pred, rpn_box, samples, matches, raw_rpn_score, raw_rpn_box, anchors, cls_targets, box_targets, box_masks, top_feat, indices)
return (cls_pred, box_pred, rpn_box, samples, matches, raw_rpn_score, raw_rpn_box, anchors, cls_targets, box_targets, box_masks, indices)
box_pred = self.box_predictor(box_feat)
box_pred = box_pred.reshape((batch_size, num_roi, self.num_class, 4))
cls_ids, scores = self.cls_decoder(F.softmax(cls_pred, axis=-1))
cls_ids = cls_ids.transpose((0, 2, 1)).reshape((0, 0, 0, 1))
scores = scores.transpose((0, 2, 1)).reshape((0, 0, 0, 1))
box_pred = box_pred.transpose((0, 2, 1, 3))
rpn_boxes = _split(rpn_box, axis=0, num_outputs=batch_size, squeeze_axis=False)
cls_ids = _split(cls_ids, axis=0, num_outputs=batch_size, squeeze_axis=True)
scores = _split(scores, axis=0, num_outputs=batch_size, squeeze_axis=True)
box_preds = _split(box_pred, axis=0, num_outputs=batch_size, squeeze_axis=True)
results = []
for rpn_box, cls_id, score, box_pred in zip(rpn_boxes, cls_ids, scores, box_preds):
bbox = self.box_decoder(box_pred, rpn_box)
res = F.concat(*[cls_id, score, bbox], dim=-1)
if self.force_nms:
res = res.reshape((1, -1, 0))
res = F.contrib.box_nms(
res, overlap_thresh=self.nms_thresh, topk=self.nms_topk, valid_thresh=0.0001,
id_index=0, score_index=1, coord_start=2, force_suppress=self.force_nms)
res = res.reshape((-3, 0))
results.append(res)
result = F.stack(*results, axis=0)
ids = F.slice_axis(result, axis=-1, begin=0, end=1)
scores = F.slice_axis(result, axis=-1, begin=1, end=2)
bboxes = F.slice_axis(result, axis=-1, begin=2, end=6)
if self._additional_output:
return ids, scores, bboxes, feat
return ids, scores, bboxes
5.模型训练
1.学习率设置
lr_steps = sorted([int(ls) for ls in lr_decay_epoch.split(',') if ls.strip()])
lr_decay_epoch = [e for e in lr_steps]
lr_scheduler = LRSequential([
LRScheduler('linear', base_lr=0, target_lr=learning_rate,
nepochs=0, iters_per_epoch=self.num_samples // self.batch_size),
LRScheduler(lr_mode, base_lr=learning_rate,
nepochs=TrainNum,
iters_per_epoch=self.num_samples // self.batch_size,
step_epoch=lr_decay_epoch,
step_factor=lr_decay, power=2),
])
2.优化器设置
if optim == 1:
trainer = gluon.Trainer(self.model.collect_params(), 'sgd', {'learning_rate': learning_rate, 'wd': 0.0005, 'momentum': 0.9, 'lr_scheduler': lr_scheduler})
elif optim == 2:
trainer = gluon.Trainer(self.model.collect_params(), 'adagrad', {'learning_rate': learning_rate, 'lr_scheduler': lr_scheduler})
else:
trainer = gluon.Trainer(self.model.collect_params(), 'adam', {'learning_rate': learning_rate, 'lr_scheduler': lr_scheduler})
3.损失设置
rpn_cls_loss = mx.gluon.loss.SigmoidBinaryCrossEntropyLoss(from_sigmoid=False)
rpn_box_loss = mx.gluon.loss.HuberLoss(rho=1. / 9.)
rcnn_cls_loss = mx.gluon.loss.SoftmaxCrossEntropyLoss()
rcnn_box_loss = mx.gluon.loss.HuberLoss(rho=1.)
metrics = [mx.metric.Loss('RPN_Conf'),
mx.metric.Loss('RPN_SmoothL1'),
mx.metric.Loss('RCNN_CrossEntropy'),
mx.metric.Loss('RCNN_SmoothL1'), ]
4.循环训练
for i, batch in enumerate(self.train_loader):
batch = self.split_and_load(batch)
metric_losses = [[] for _ in metrics]
add_losses = [[] for _ in metrics2]
if executor is not None:
for data in zip(*batch):
executor.put(data)
for j in range(len(self.ctx)):
if executor is not None:
result = executor.get()
else:
result = rcnn_task.forward_backward(list(zip(*batch))[0])
for k in range(len(metric_losses)):
metric_losses[k].append(result[k])
for k in range(len(add_losses)):
add_losses[k].append(result[len(metric_losses) + k])
trainer.step(self.batch_size)
for metric, record in zip(metrics, metric_losses):
metric.update(0, record)
for metric, records in zip(metrics2, add_losses):
for pred in records:
metric.update(pred[0], pred[1])
msg = ','.join(['{}={:.3f}'.format(*metric.get()) for metric in metrics + metrics2])
print('[Epoch {}][Batch {}], Speed: {:.3f} samples/sec, {}'.format(epoch, i, self.batch_size / (time.time() - btic), msg))
btic = time.time()
6.模型预测
def predict(self, image, confidence=0.5, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
start_time = time.time()
origin_img = copy.deepcopy(image)
base_imageSize = origin_img.shape
image = cv2.cvtColor(image,cv2.COLOR_RGB2BGR)
image = cv2.resize(image,(self.min_size,self.min_size))
img = nd.array(image)
img = mx.nd.image.to_tensor(img)
img = mx.nd.image.normalize(img, mean=mean, std=std)
x = img.expand_dims(0)
x = x.as_in_context(self.ctx[0])
labels, scores, bboxes = [xx[0].asnumpy() for xx in self.model(x)]
origin_img_pillow = self.cv2_pillow(origin_img)
font = ImageFont.truetype(font='./model_data/simhei.ttf', size=np.floor(3e-2 * np.shape(origin_img_pillow)[1] + 0.5).astype('int32'))
thickness = max((np.shape(origin_img_pillow)[0] + np.shape(origin_img_pillow)[1]) // self.min_size, 1)
imgbox = []
for i, bbox in enumerate(bboxes):
if (scores is not None and scores.flat[i] < confidence) or labels is not None and labels.flat[i] < 0:
continue
cls_id = int(labels.flat[i]) if labels is not None else -1
xmin, ymin, xmax, ymax = [int(x) for x in bbox]
xmin = int(xmin / self.min_size * base_imageSize[1])
xmax = int(xmax / self.min_size * base_imageSize[1])
ymin = int(ymin / self.min_size * base_imageSize[0])
ymax = int(ymax / self.min_size * base_imageSize[0])
# print(xmin, ymin, xmax, ymax, self.classes_names[cls_id])
class_name = self.classes_names[cls_id]
score = '{:d}%'.format(int(scores.flat[i] * 100)) if scores is not None else ''
imgbox.append([(xmin, ymin, xmax, ymax), cls_id, self.classes_names[cls_id], score])
top, left, bottom, right = ymin, xmin, ymax, xmax
label = '{}-{}'.format(class_name, score)
draw = ImageDraw.Draw(origin_img_pillow)
label_size = draw.textsize(label, font)
label = label.encode('utf-8')
if top - label_size[1] >= 0:
text_origin = np.array([left, top - label_size[1]])
else:
text_origin = np.array([left, top + 1])
for i in range(thickness):
draw.rectangle([left + i, top + i, right - i, bottom - i], outline=self.colors[cls_id])
draw.rectangle([tuple(text_origin), tuple(text_origin + label_size)], fill=self.colors[cls_id])
draw.text(text_origin, str(label,'UTF-8'), fill=(0, 0, 0), font=font)
del draw
result_value = {
"image_result": self.pillow_cv2(origin_img_pillow),
"bbox": imgbox,
"time": (time.time() - start_time) * 1000
}
return result_value
四、算法主入口
if __name__ == '__main__':
# ctu = Ctu_FasterRcnn(USEGPU='0',ampFlag=True,mixupFlag=False,min_size = 416,max_size=416)
# ctu.InitModel(DataDir=r'D:/Ctu/Ctu_Project_DL/DataSet/DataSet_Detection_YaoPian',batch_size=1,num_workers = 0,model_name='faster_rcnn_resnet18_v1b_voc',pre_Model='./Model_faster_rcnn_resnet18_v1b_voc/best_model.dat')
# ctu.train(TrainNum=300,learning_rate=0.0001,lr_decay_epoch='30,70,150,200',ModelPath='./Model',lr_decay=0.9,disable_hybridization=False)
ctu = Ctu_FasterRcnn(USEGPU='0',ampFlag=True)
ctu.LoadModel(r'./Model_faster_rcnn_resnet50_v1b_voc')
cv2.namedWindow("result", 0)
cv2.resizeWindow("result", 640, 480)
index = 0
for root, dirs, files in os.walk(r'D:/Ctu/Ctu_Project_DL\DataSet\DataSet_Detection_Color/test'):
for f in files:
img_cv = ctu.read_image(os.path.join(root, f))
if img_cv is None:
continue
res = ctu.predict(img_cv, 0.7)
for each in res['bbox']:
print(each)
print("耗时:" + str(res['time']) + ' ms')
# cv2.imwrite(str(index + 1)+'.bmp',res['image_result'])
cv2.imshow("result", res['image_result'])
cv2.waitKey()
# index +=1
五、训练效果展示
备注：项目模型的本人没有保存因此会后续提供训练效果
3976
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
