1.数据集下载
http://cg.cs.tsinghua.edu.cn/traffic-sign/
2.Yolov4模型下载
- README.md
https://gitee.com/xtlower/yolov4-pytorch?_from=gitee_search
3.数据集分析1
# 3.1 格式化文件:format_json.py
import json
with open('./instances_val2017.json', 'rb') as fp:
data = json.load(fp)
print(type(data))
print(data.keys())
with open('./instances_val2017_formatted.json', 'w') as fp:
json.dump(data, fp, indent=4, sort_keys=True, ensure_ascii=False)
# 3.2 jupyter notebook 数据集分析
import json
with open('./annotations_all_formatted.json', 'rb') as fp: # 读文件
obj = json.load(fp)
import jsonpath # 统计每个类别出现的次数
cats = jsonpath.jsonpath(obj, '$..category') # 把obj中所有category拿到
len(cats)
27346
len(obj['types'])
232
from collections import Counter
c = Counter(cats) # 统计cats中每个类别出现的次数可用Counter
c
Counter({'ph2': 12,
'pl40': 1413,
...
'w62': 1,
'w44': 1})
len(c)
201
# 3.3 画图显示类别分布
import matplotlib.pyplot as plt
import numpy as np
plt.figure(figsize=(7, 40))
plt.barh(np.arange(len(c.keys())), c.values())
_ = plt.yticks(np.arange(len(c.keys())), c.keys())
4.数据集预处理
new_c = dict((k, v) for k, v in c.items() if v >= 5) # 取出最少出现5次类别, 用来参与训练
new_c.keys()
dict_keys(['ph2', 'pl40', 'pl30', 'p26', 'p27', 'i5', 'pne', 'pr60', 'ph4.5', 'pr20', 'pbp', 'ip', 'p5', 'pn', 'p11', 'w30', 'p9', 'pm20', 'pl5', 'i4', 'pl70', 'pl80', 'il60', 'il70', 'il80', 'ph4', 'pl100', 'pl50', 'i2', 'im', 'pdd', 'pg', 'w59', 'p3', 'pm30', 'p1', 'pl60', 'p12', 'p17', 'p18', 'pr50', 'pn-2', 'il110', 'il90', 'pl120', 'w57', 'p10', ...
'pl10', 'phclr', 'i12', 'pm2', 'pl90', 'p14', 'pm8', 'pm40', 'p2', 'i3', 'p25', 'pm25', 'p15', 'w8', 'w24', 'w41', 'pmr', 'i1', 'w20', 'pr80', 'w10', 'i14', 'ph4.8', 'phcs', 'pctl', 'pt', 'ph4.3', 'w18', 'i15'])
5.数据集转换
- tt100k2coco.py
import json
import os
import cv2
_type = 'test'
src = '../../tt100k_2021'
anno = os.path.join(src, 'annotations_all.json')
dataset = {'categories': [], 'images': [], 'annotations': []}
categories = ['__background__']
with open(anno, 'r') as f:
annotations = json.load(f)
cat = ['ph2', 'pl40', 'pl30', 'p26', 'p27', 'i5', 'pne', 'pr60', 'ph4.5', 'pr20', 'pbp', 'ip', 'p5', 'pn', 'p11', 'w30', 'p9', 'pm20', 'pl5', 'i4', 'pl70', 'pl80', 'il60', 'il70', 'il80', 'ph4', 'pl100', 'pl50', 'i2', 'im', 'pdd', 'pg', 'w59', 'p3', 'pm30', 'p1', 'pl60', 'p12', 'p17', 'p18', 'pr50', 'pn-2', 'il110', 'il90', 'pl120', 'w57', 'p10', 'pm15', 'w55', 'i2r', 'p13', 'pl20', 'w13', 'i4l', 'pm55', 'w63', 'pss', 'pcr', 'p16', 'w32', 'pb', 'w21', 'il100', 'ph5', 'pl110', 'w58', 'pc', 'p23', 'p19', 'pbm', 'ps', 'pl25', 'pl35', 'pr40', 'w45', 'ph2.5', 'p6', 'pcl', 'w42', 'ph3.5', 'ph3.3', 'w22', 'w43', 'pm50', 'pm5', 'ph4.2', 'p29', 'pcd', 'w3', 'pr70', 'p8', 'pl15', 'pa13', 'ph2.8', 'pa14', 'wc', 'ph3', 'w47', 'pm10', 'il50', 'pr30', 'iz', 'w46', 'ph2.2', 'pm35', 'p1n', 'i10', 'w26', 'w34', 'pw3.2', 'pw4', 'pa10', 'pm46', 'pmb', 'i13', 'pl10', 'phclr', 'i12', 'pm2', 'pl90', 'p14', 'pm8', 'pm40', 'p2', 'i3', 'p25', 'pm25', 'p15', 'w8', 'w24', 'w41', 'pmr', 'i1', 'w20', 'pr80', 'w10', 'i14', 'ph4.8', 'phcs', 'pctl', 'pt', 'ph4.3', 'w18', 'i15']
for i, c in enumerate(cat, 1):
dataset['categories'].append(
{'id': i, 'name': c, 'supercategory': c[0]})
categories.append(c)
imgs = annotations['imgs'] # imgs本身是个字典
types = annotations['types']
g_index = 1
for idx, (img, info) in enumerate(imgs.items(), 1):
if _type in info['path']:
bgr = cv2.imread(os.path.join(src, info['path']))
h, w, _ = bgr.shape
objects = info['objects']
dataset['images'].append({
'file_name': info['path'].split('/')[-1],
'id': idx,
'width': w,
'height': h})
for obj in objects:
bbox = obj['bbox']
x1 = bbox['xmin']
y1 = bbox['ymin']
x2 = bbox['xmax']
y2 = bbox['ymax']
label = obj['category']
segmentation = [[x1, y1, x2, y1, x2, y2, x1, y2]]
width = max(0, x2 - x1)
height = max(0, y2 - y1)
if label not in categories:
continue
dataset['annotations'].append({
'area': width * height,
'bbox': [x1, y1, width, height],
'category_id': categories.index(label),
'id': g_index,
'image_id': idx,
'iscrowd': 0,
'segmentation': segmentation})
g_index += 1
json_name = os.path.join(src, f'TT100K_CoCo_format_{_type}.json')
with open(json_name, 'w') as f:
json.dump(dataset, f)
print('done')
6.数据集分析 2
- coco工具包的使用
from pycocotools.coco import COCO
json_path = '../tt100k_2021/annotations/TT100K_CoCo_format_train.json'
coco = COCO(annotation_file=json_path)
loading annotations into memory...
Done (t=0.12s)
creating index...
index created!
coco.anns # 注释信息
{1: {'area': 3519.984399989997,
'bbox': [1580.0, 758.667, 58.66669999999999, 59.99969999999996],
'category_id': 64,
'id': 1,
'image_id': 1,
'iscrowd': 0,
'segmentation': [[1580.0,
758.667,
1638.6667,
758.667,
1638.6667,
818.6667,
1580.0,
818.6667]]},
...}
coco.imgs # 图片信息
{1: {'file_name': '62627.jpg', 'id': 1, 'width': 2048, 'height': 2048},
...}
coco.cats # 类别信息
{1: {'id': 1, 'name': 'ph2', 'supercategory': 'p'},
...
144: {'id': 144, 'name': 'i15', 'supercategory': 'i'}}
ids = list(sorted(coco.imgs.keys())) # 图片id
len(ids) # 图片数量
6034
coco.loadImgs([1, 3])
[{'file_name': '62627.jpg', 'id': 1, 'width': 2048, 'height': 2048},
{'file_name': '84351.jpg', 'id': 3, 'width': 2048, 'height': 2048}]
img_path = '../tt100k_2021/train'
from PIL import Image, ImageDraw
import os
coco_classes = dict((v['id'], v['name']) for k,v in coco.cats.items())
import matplotlib.pyplot as plt
imgs = []
for img_id in ids[:3]: # 取前3张图给大家展示
ann_ids = coco.getAnnIds(imgIds=img_id) # 获取图片对应id的所有annotations的id信息
targets = coco.loadAnns(ann_ids) # 根据annotations idx的信息获取所有标注信息
path = coco.loadImgs(img_id)[0]['file_name'] # 获取文件名
img = Image.open(os.path.join(img_path, path)).convert('RGB') # 读图片
draw = ImageDraw.Draw(img)
for target in targets: # 画出目标框
x, y, w, h = target['bbox']
x1, y1, x2, y2 = x, y, int(x + w), int(y + h)
draw.rectangle((x1, y1, x2, y2))
draw.text((x1, y1), coco_classes[target['category_id']])
imgs.append(img)
# plt.figure(figsize=(20, 20))
# plt.imshow(img)
# plt.show()
imgs[0]
imgs[1]
coco_classes # 数据类别
{1: 'ph2',
2: 'pl40',
...
143: 'w18',
144: 'i15'}
s = '\n'.join(list(coco_classes.values())) # 数据类别文件
with open('./model_data/tt100k_144_classes.txt', 'w') as f:
f.write(s)
7. 训练
- train.py文件
import os
# os.environ["CUDA_VISIBLE_DEVICES"] = "2"
import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.optim as optim
from torch.utils.data import DataLoader
from tqdm import tqdm
from nets.yolo4 import YoloBody
from nets.yolo_training import LossHistory, YOLOLoss, weights_init
# from utils.dataloader import YoloDataset, yolo_dataset_collate
from utils.coco import COCO, yolo_dataset_collate, COCOEval
from utils.utils import DecodeBox, non_max_suppression, add_weight_decay, setup_seed
from utils.summary import Summary
import shutil
yolo_decodes = []
confidence = 0.01
iou = 0.5
g_steps = 0
def get_classes(classes_path): # 获得类和先验框
with open(classes_path) as f: # 获得类
class_names = f.readlines()
class_names = [c.strip() for c in class_names]
return class_names
def get_anchors(anchors_path):
with open(anchors_path) as f: # 文件加载先验框
anchors = f.readline()
anchors = [float(x) for x in anchors.split(',')]
return np.array(anchors).reshape([-1, 3, 2])[::-1, :, :]
def get_lr(optimizer):
for param_group in optimizer.param_groups:
return param_group['lr']
def fit_one_epoch(net, yolo_loss, epoch, epoch_size, epoch_size_val, gen, genval, Epoch, cuda):
total_loss = 0
val_loss = 0
net.train()
print('Start Train')
with tqdm(total=epoch_size, desc=f'Epoch {epoch + 1}/{Epoch}', postfix=dict, mininterval=0.3) as pbar:
for iteration, batch in enumerate(gen):
if iteration >= epoch_size:
break
# images, targets = batch[0], batch[1]
img_ids, images, targets = batch[0], batch[1], batch[2]
with torch.no_grad():
if cuda:
images = torch.from_numpy(images).type(torch.FloatTensor).cuda()
targets = [torch.from_numpy(ann).type(torch.FloatTensor) for ann in targets]
else:
images = torch.from_numpy(images).type(torch.FloatTensor)
targets = [torch.from_numpy(ann).type(torch.FloatTensor) for ann in targets]
optimizer.zero_grad()
outputs = net(images)
losses = []
num_pos_all = 0
for i in range(3): # 计算损失
loss_item, num_pos = yolo_loss(outputs[i], targets)
losses.append(loss_item)
num_pos_all += num_pos
loss = sum(losses) / num_pos_all
total_loss += loss.item()
loss.backward()
optimizer.step()
summary.add_scalar('train/total_loss', total_loss / (iteration + 1), epoch * epoch_size + iteration)
pbar.set_postfix(**{'total_loss': total_loss / (iteration + 1),
'lr': get_lr(optimizer)})
pbar.update(1)
net.eval()
print('Start Validation')
with torch.no_grad():
img_ids = []
detections = []
with tqdm(total=epoch_size_val, desc=f'Epoch {epoch + 1}/{Epoch}', postfix=dict, mininterval=0.3) as pbar:
for iteration, batch in enumerate(genval):
if iteration >= epoch_size_val:
break
# images_val, targets_val = batch[0], batch[1]
img_ids_val, images_val, targets_val = batch[0], batch[1], batch[2]
with torch.no_grad():
if cuda:
images_val = torch.from_numpy(images_val).type(torch.FloatTensor).cuda()
targets_val = [torch.from_numpy(ann).type(torch.FloatTensor) for ann in targets_val]
else:
images_val = torch.from_numpy(images_val).type(torch.FloatTensor)
targets_val = [torch.from_numpy(ann).type(torch.FloatTensor) for ann in targets_val]
optimizer.zero_grad()
outputs = net(images_val)
output_list = [] # 计算mAP
for ii in range(3):
output_list.append(yolo_decodes[ii](outputs[ii]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output, num_classes, conf_thres=confidence,
nms_thres=iou)
img_ids += img_ids_val
detections += batch_detections
losses = []
num_pos_all = 0
for i in range(3):
loss_item, num_pos = yolo_loss(outputs[i], targets_val)
losses.append(loss_item)
num_pos_all += num_pos
loss = sum(losses) / num_pos_all
val_loss += loss.item()
# 将loss写入tensorboard, 下面注释的是每一步都写
# if Tensorboard:
# writer.add_scalar('Val_loss', loss, val_tensorboard_step)
# val_tensorboard_step += 1
summary.add_scalar('val/total_loss', val_loss / (iteration + 1), epoch * epoch_size + iteration)
pbar.set_postfix(**{'total_loss': val_loss / (iteration + 1)})
pbar.update(1)
loss_history.append_loss(total_loss / (epoch_size + 1), val_loss / (epoch_size_val + 1))
eval_results = val_dataset.run_eval(img_ids, detections)
summary.add_scalar('val/mAP', eval_results[0], epoch + 1)
print('Finish Validation')
print('Epoch:' + str(epoch + 1) + '/' + str(Epoch))
print(eval_results)
print('Total Loss: %.4f || Val Loss: %.4f ' % (total_loss / (epoch_size + 1), val_loss / (epoch_size_val + 1)))
net_save_path = os.path.join(os.path.join(exp_dir, "ckpt"), "model_latest.pth")
net_save_path_best = os.path.join(os.path.join(exp_dir, "ckpt"), "best_weights.pth")
save_best = False
if eval_results[0] >= metrics['mAP']:
save_best = True
metrics['train_loss'] = total_loss / (epoch_size + 1)
metrics['val_loss'] = val_loss / (epoch_size_val + 1)
metrics['mAP'] = eval_results[0]
metrics['best_model_epoch'] = epoch + 1
print('Saving state, iter:', str(epoch + 1))
torch.save({"state_dict": model.state_dict(),
"metric": metrics}
, net_save_path)
if save_best:
shutil.copy(net_save_path, net_save_path_best)
print("Saving current best: {}, metric:{}".format(net_save_path_best, metrics))
print("@@@ best metric:{}".format(net_save_path_best, metrics)
if __name__ == "__main__":
setup_seed(1215)
Cuda = True # 是否使用Cuda 没有GPU可以设置成False
# ------------------------------------------------------#
# 是否对损失进行归一化,用于改变loss的大小
# 用于决定计算最终loss是除上batch_size还是除上正样本数量
# ------------------------------------------------------#
normalize = False
# -------------------------------#
# 输入的shape大小
# 显存比较小可以使用416x416
# 显存比较大可以使用608x608
# -------------------------------#
input_shape = (416, 416)
# input_shape = (608, 608)
# ----------------------------------------------------#
# classes和anchor的路径,非常重要
# 训练前一定要修改classes_path,使其对应自己的数据集
# ----------------------------------------------------#
anchors_path = 'model_data/yolo_anchors.txt'
classes_path = 'model_data/tt100k_144_classes.txt'
# ------------------------------------------------------#
# Yolov4的tricks应用
# mosaic 马赛克数据增强 True or False
# 实际测试时mosaic数据增强并不稳定,所以默认为False
# Cosine_scheduler 余弦退火学习率 True or False
# label_smoothing 标签平滑 0.01以下一般 如0.01、0.005
# ------------------------------------------------------#
mosaic = False
Cosine_lr = True
smoooth_label = 0.005
weight_decay = 1e-5
metrics = {'mAP': 0, 'train_loss': float('inf'), 'val_loss': float('inf'), 'best_model_epoch': 0}
class_names = get_classes(classes_path) # 获取classes和anchor
anchors = get_anchors(anchors_path)
num_classes = len(class_names)
for i in range(3):
yolo_decodes.append(
DecodeBox(anchors[i], num_classes, (input_shape[1], input_shape[0])))
# ------------------------------------------------------#
# 创建yolo模型 训练前一定要修改classes_path和对应的txt文件
# ------------------------------------------------------#
model = YoloBody(len(anchors[0]), num_classes)
weights_init(model)
# 想从头训练, 不使用预训练的权重, 就下面这段代码注释掉
model_path = "model_data/yolo4_weights.pth"
# model_path = "exp/exp_baseline/ckpt/best_weights.pth"
print('Loading weights into state dict...')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)
model_dict = model.state_dict()
pretrained_dict = torch.load(model_path, map_location=device)["state_dict"]
pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
print('Finished!')
net = model.train()
if Cuda:
net = torch.nn.DataParallel(model)
cudnn.benchmark = True
net = net.cuda()
yolo_loss = YOLOLoss(np.reshape(anchors, [-1, 2]), num_classes, (input_shape[1], input_shape[0]), smoooth_label,Cuda, normalize)
exp_dir = "./exp/exp_baseline_mosaic"
loss_history = LossHistory(exp_dir)
summary = Summary(os.path.join(exp_dir, "summary"))
ckpt = os.path.join(exp_dir, "ckpt")
if not os.path.exists(ckpt):
os.makedirs(ckpt)
if Cuda:
graph_inputs = torch.randn(1, 3, input_shape[0], input_shape[1]).type(torch.FloatTensor).cuda()
else:
graph_inputs = torch.randn(1, 3, input_shape[0], input_shape[1]).type(torch.FloatTensor)
summary.add_graph(model, graph_inputs)
data_dir = "../tt100k_2021/" # 获得图片路径和标签
# ------------------------------------------------------#
# 主干特征提取网络特征通用,冻结训练可以加快训练速度
# 也可以在训练初期防止权值被破坏。
# Init_Epoch为起始世代
# Freeze_Epoch为冻结训练的世代
# Epoch总训练世代
# 提示OOM或者显存不足请调小Batch_size
# ------------------------------------------------------#
if True:
lr = 1e-3
Batch_size = 16
Init_Epoch = 0
Freeze_Epoch = 30
# ----------------------------------------------------------------------------#
# YOLOv5代码中conv层和FC层的bias参数,以及BN层参数并不进行权重衰减,此处采用这个方法
# ----------------------------------------------------------------------------#
parameters = add_weight_decay(net, weight_decay)
optimizer = optim.Adam(parameters, lr)
# optimizer = optim.Adam(net.parameters(), lr)
if Cosine_lr:
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=5, eta_min=1e-5)
else:
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.92)
train_dataset = COCO(data_dir, (input_shape[0], input_shape[1]), mosaic=mosaic)
val_dataset = COCOEval(data_dir, (input_shape[0], input_shape[1]))
# train_dataset = YoloDataset(lines[:num_train], (input_shape[0], input_shape[1]), mosaic=mosaic, is_train=True)
# val_dataset = YoloDataset(lines[num_train:], (input_shape[0], input_shape[1]), mosaic=False, is_train=False)
gen = DataLoader(train_dataset, shuffle=True, batch_size=Batch_size, num_workers=4, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
gen_val = DataLoader(val_dataset, shuffle=True, batch_size=Batch_size, num_workers=4, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
epoch_size = len(train_dataset) // Batch_size
epoch_size_val = len(val_dataset) // Batch_size
if epoch_size == 0 or epoch_size_val == 0:
raise ValueError("数据集过小,无法进行训练,请扩充数据集。")
for param in model.backbone.parameters(): # 冻结一定部分训练
param.requires_grad = False
for epoch in range(Init_Epoch, Freeze_Epoch):
fit_one_epoch(net, yolo_loss, epoch, epoch_size, epoch_size_val, gen, gen_val, Freeze_Epoch, Cuda)
lr_scheduler.step()
torch.cuda.empty_cache() # 释放所有未占用的缓存内存
if True:
lr = 1e-4
Batch_size = 8
Freeze_Epoch = 30
Unfreeze_Epoch = 100
# ----------------------------------------------------------------------------#
# 实际测试时,发现optimizer的weight_decay起到了反作用,
# 所以去除掉了weight_decay,大家也可以开起来试试,一般是weight_decay=5e-4
# ----------------------------------------------------------------------------#
parameters = add_weight_decay(net, weight_decay)
optimizer = optim.Adam(parameters, lr)
# optimizer = optim.Adam(net.parameters(), lr)
if Cosine_lr:
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=5, eta_min=1e-5)
else:
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.92)
train_dataset = COCO(data_dir, (input_shape[0], input_shape[1]), mosaic=mosaic)
val_dataset = COCOEval(data_dir, (input_shape[0], input_shape[1]))
# train_dataset = YoloDataset(lines[:num_train], (input_shape[0], input_shape[1]), mosaic=mosaic, is_train=True)
# val_dataset = YoloDataset(lines[num_train:], (input_shape[0], input_shape[1]), mosaic=False, is_train=False)
gen = DataLoader(train_dataset, shuffle=True, batch_size=Batch_size, num_workers=4, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
gen_val = DataLoader(val_dataset, shuffle=True, batch_size=Batch_size, num_workers=4, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
epoch_size = len(train_dataset) // Batch_size
epoch_size_val = len(val_dataset) // Batch_size
if epoch_size == 0 or epoch_size_val == 0:
raise ValueError("数据集过小,无法进行训练,请扩充数据集。")
for param in model.backbone.parameters(): # 解冻后训练
param.requires_grad = True
for epoch in range(Freeze_Epoch, Unfreeze_Epoch):
fit_one_epoch(net, yolo_loss, epoch, epoch_size, epoch_size_val, gen, gen_val, Unfreeze_Epoch, Cuda)
lr_scheduler.step()
8.预测
- predict.py文件
# ----------------------------------------------------#
# 对视频中的predict.py进行了修改,将单张图片预测、摄像头检测和FPS测试功能
# 整合到了一个py文件中,通过指定mode进行模式的修改
# ----------------------------------------------------#
import time
import cv2
import numpy as np
from PIL import Image
from yolo import YOLO
if __name__ == "__main__":
yolo = YOLO()
# -------------------------------------------------------------------------#
# mode用于指定测试的模式
# 'predict'表示单张图片预测
# 'video'表示视频检测
# 'fps'表示测试fps
# -------------------------------------------------------------------------#
mode = "video"
# -------------------------------------------------------------------------#
# video_path用于指定视频的路径,当video_path=0时表示检测摄像头
# video_save_path表示视频保存的路径,当video_save_path=""时表示不保存
# video_fps用于保存的视频的fps
# video_path、video_save_path和video_fps仅在mode='video'时有效
# 保存视频时需要ctrl+c退出才会完成完整的保存步骤,不可直接结束程序
# -------------------------------------------------------------------------#
video_path = './video.mp4'
video_save_path = ""
video_fps = 25.0
if mode == "predict":
'''
1、该代码无法直接进行批量预测,如果想要批量预测,可以利用os.listdir()遍历文件夹,利用Image.open打开图片文件进行预测。具体流程可以参考get_dr_txt.py,在get_dr_txt.py即实现了遍历还实现了目标信息的保存。
2、如果想要进行检测完的图片的保存,利用r_image.save("img.jpg")即可保存,直接在predict.py里进行修改即可。
3、如果想要获得预测框的坐标,可以进入yolo.detect_image函数,在绘图部分读取top,left,bottom,right这四个值。
4、如果想要利用预测框截取下目标,可以进入yolo.detect_image函数,在绘图部分利用获取到的top,left,bottom,right这四个值在原图上利用矩阵的方式进行截取。
5、如果想要在预测图上写额外的字,比如检测到的特定目标的数量,可以进入yolo.detect_image函数,在绘图部分对predicted_class进行判断,比如判断if predicted_class == 'car': 即可判断当前目标是否为车,然后记录数量即可。利用draw.text即可写字。
'''
while True:
img = input('Input image filename:')
try:
image = Image.open(img)
except:
print('Open Error! Try again!')
continue
else:
r_image = yolo.detect_image(image)
r_image.save(img.split("/")[-1])
r_image.show()
elif mode == "video":
capture = cv2.VideoCapture(video_path)
if video_save_path != "":
fourcc = cv2.VideoWriter_fourcc(*'XVID')
size = (int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))
out = cv2.VideoWriter(video_save_path, fourcc, video_fps, size)
fps = 0.0
while (True):
t1 = time.time()
ref, frame = capture.read() # 读取某一帧
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # 格式转变BGRtoRGB
frame = Image.fromarray(np.uint8(frame)) # 转变成Image
frame = np.array(yolo.detect_image(frame)) # 进行检测
frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) # RGBtoBGR满足opencv显示格式
fps = (fps + (1. / (time.time() - t1))) / 2
print("fps= %.2f" % (fps))
frame = cv2.putText(frame, "fps= %.2f" % (fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
cv2.imshow("video", frame)
c = cv2.waitKey(1) & 0xff
if video_save_path != "":
out.write(frame)
if c == 27:
capture.release()
break
capture.release()
out.release()
cv2.destroyAllWindows()
elif mode == "fps":
test_interval = 100
img = Image.open('img/street.jpg')
tact_time = yolo.get_FPS(img, test_interval)
print(str(tact_time) + ' seconds, ' + str(1 / tact_time) + 'FPS, @batch_size 1')
else:
raise AssertionError("Please specify the correct mode: 'predict', 'video' or 'fps'.")