Pytorch版 Yolo v4:
https://github.com/Tianxiaomo/pytorch-YOLOv4
COCO数据集预训练模型
yolov4.pth(https://pan.baidu.com/s/1ZroDvoGScDgtE1ja_QqJVw Extraction code:xrq9)
yolov4.conv.137.pth(https://pan.baidu.com/s/1ovBie4YyVQQoUrC3AY0joA Extraction code:kcel)
配置环境
conda create -n yolov4
安装requirements.txt依赖库
模型推断
python models.py <num_classes> <weightfile> <imgfile> <IN_IMAGE_H> <IN_IMAGE_W> <namefile(optional)>
修改代码,多图片预测
1 准备图片路径名
文本保存为train.txt,每行一个文件名
2 修改测试代码
class Yolov4(nn.Module):
def __init__(self, yolov4conv137weight=None, n_classes=80, inference=False):
super().__init__()
output_ch = (4 + 1 + n_classes) * 3
# backbone
self.down1 = DownSample1()
self.down2 = DownSample2()
self.down3 = DownSample3()
self.down4 = DownSample4()
self.down5 = DownSample5()
# neck
self.neek = Neck(inference)
# yolov4conv137
if yolov4conv137weight:
_model = nn.Sequential(self.down1, self.down2, self.down3, self.down4, self.down5, self.neek)
pretrained_dict = torch.load(yolov4conv137weight)
model_dict = _model.state_dict()
# 1. filter out unnecessary keys
pretrained_dict = {k1: v for (k, v), k1 in zip(pretrained_dict.items(), model_dict)}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
_model.load_state_dict(model_dict)
# head
self.head = Yolov4Head(output_ch, n_classes, inference)
def forward(self, input):
d1 = self.down1(input)
d2 = self.down2(d1)
d3 = self.down3(d2)
d4 = self.down4(d3)
d5 = self.down5(d4)
x20, x13, x6 = self.neek(d5, d4, d3)
output = self.head(x20, x13, x6)
return output
if __name__ == "__main__":
import sys
import cv2
import pandas as pd
from tqdm import tqdm
import pickle
weightfile = 'ckpt/yolov4.pth'
train_data_path = 'data/small_data_list/small_data_0.2_test_data_list.txt'
root_dir = '/home/flyingbird/Data/images'
conv137weight = 'ckpt/yolov4.conv.137.pth'
n_classes = 80
width = 512
height = 512
all_box_dict = {}
model = Yolov4(yolov4conv137weight=conv137weight, n_classes=n_classes, inference=True)
pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda'))
model.load_state_dict(pretrained_dict)
use_cuda = True
if use_cuda:
model.cuda()
test_data_list = list(pd.read_csv(train_data_path, index_col=False, header=None, sep=' ')[0])
for i in tqdm(test_data_list):
img_path = os.path.join(root_dir, str(i)) + '.jpg'
img = cv2.imread(img_path)
# Inference input size is 416*416 does not mean training size is the same
# Training size could be 608*608 or even other sizes
# Optional inference sizes:
# Hight in {320, 416, 512, 608, ... 320 + 96 * n}
# Width in {320, 416, 512, 608, ... 320 + 96 * m}
sized = cv2.resize(img, (width, height))
sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
from tool.utils import load_class_names, plot_boxes_cv2
from tool.torch_utils import do_detect
for num in range(2):
# This 'for' loop is for speed check
# Because the first iteration is usually longer
boxes = do_detect(model, sized, 0.00001, 0.1, use_cuda)
# all_box_dict[i] = boxes[0]
# print(len(boxes[0]))
namesfile = 'data/coco.names'
class_names = load_class_names(namesfile)
save_path = os.path.join('predict/test', str(i)) + '.jpg'
img, boxes_list = plot_boxes_cv2(img, boxes[0], save_path, class_names)
# print('red one')
all_box_dict[i] = boxes_list
with open('predict/test_pred_box_dict_conf_0.00001_nms_0.1.txt', 'wb') as f:
pickle.dump(all_box_dict, f)
3 结果输出保存
前面网络结构部分不用修改,后面测试时读取每张图,检测并输出box和置信度,筛选得到最终预测结果并保存在save_path中,预测结果的box用pickle保存为txt文件.