YOLOv5二阶段目标检测

有吸烟、打电话等小目标时，可以先检测出人再检测人框里是否有吸烟打电话。该脚本用来测试二阶段精度。

直接 pip install yolov5

import yolov5
from PIL import Image
import os
import shutil
import glob
# 载入模型
model = yolov5.load(r'E:\yolov5-master\yolofuwuqi\person.pt')
model2 = yolov5.load(r'E:\yolov5-master\runs\train\exp20\weights\best.pt')
if not os.path.exists('results'):
    os.mkdir('results')
if not os.path.exists('output'):
    os.mkdir('output')
# 使用特定尺寸进行推理
# results = model(img, size=1280)
img_path = r"E:\yolov5-master\some_demo\222"
for img0 in os.listdir(img_path):
    if 'jpg' not in img0:
        continue
    print(img0)
    img = os.path.join(img_path,img0)
    im =Image.open(img)
    # 数据增强，能够检测出更多的目标，当然也有可能出现误检
    results = model(img, augment=True)
    # 检测结果数据解析，所属类别、置信度、目标位置信息
    predictions = results.pred[0]
    boxes = predictions[:, :4]   # x1, x2, y1, y2
    for box in boxes:
        im1 = im.crop([int(box[0]) , int(box[1]), int(box[2]), int(box[3])])
        results1 = model2(im1, augment=True)
        # 保存检测结果图片
        results1.save(save_dir='results/img')


for j,i in enumerate(glob.glob('results\*\*jpg')):
    shutil.move(i,'output/'+str(j)+'.jpg')

os.remove('results')