【目标检测】YOLOv5算法实现(九):模型预测

本系列文章记录本人硕士阶段YOLO系列目标检测算法自学及其代码实现的过程。其中算法具体实现借鉴于ultralytics YOLO源码Github,删减了源码中部分内容,满足个人科研需求。

本系列文章主要以YOLOv5为例完成算法的实现,后续修改、增加相关模块即可实现其他版本的YOLO算法。

文章地址:
YOLOv5算法实现(一):算法框架概述
YOLOv5算法实现(二):模型加载
YOLOv5算法实现(三):数据集加载
YOLOv5算法实现(四):正样本匹配与损失计算
YOLOv5算法实现(五):预测结果后处理
YOLOv5算法实现(六):评价指标及实现
YOLOv5算法实现(七):模型训练
YOLOv5算法实现(八):模型验证
YOLOv5算法实现(九):模型预测

本文目录

引言

本篇文章综合之前文章中的功能,实现模型的预测。模型预测的逻辑如图1所示。


图1 模型预测流程

模型预测(predict.py)

python 复制代码
def predice():
    img_size = 640  # 必须是32的整数倍 [416, 512, 608]
    file = "yolov5s"
    cfg = f"cfg/models/{file}.yaml"  # 改成生成的.cfg文件
    weights_path = f"weights/{file}/{file}.pt"  # 改成自己训练好的权重文件
    json_path = "data/dataset.json"  # json标签文件
    img_path = "test.jpg"
    save_path = f"results/{file}/test_result8.jpg"
    assert os.path.exists(cfg), "cfg file {} dose not exist.".format(cfg)
    assert os.path.exists(weights_path), "weights file {} dose not exist.".format(weights_path)
    assert os.path.exists(json_path), "json file {} dose not exist.".format(json_path)
    assert os.path.exists(img_path), "image file {} dose not exist.".format(img_path)

    with open(json_path, 'r') as f:
        class_dict = json.load(f)

    category_index = {str(v): str(k) for k, v in class_dict.items()}

    input_size = (img_size, img_size)

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    device = torch.device("cpu")
    model = Model(cfg, ch=3, nc=3)
    weights_dict = torch.load(weights_path, map_location='cpu')
    weights_dict = weights_dict["model"] if "model" in weights_dict else weights_dict
    model.load_state_dict(weights_dict, strict=False)
    model.to(device)

    model.eval()
    with torch.no_grad():
        # init
        img = torch.zeros((1, 3, img_size, img_size), device=device)
        model(img)

        img_o = cv2.imread(img_path)  # BGR
        assert img_o is not None, "Image Not Found " + img_path

        img = letterbox(img_o, new_shape=input_size, auto=True, color=(0, 0, 0))[0]

        # Convert
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
        img = np.ascontiguousarray(img)

        img = torch.from_numpy(img).to(device).float()
        img /= 255.0  # scale (0, 255) to (0, 1)
        img = img.unsqueeze(0)  # add batch dimension

        t1 = torch_utils.time_synchronized()
        pred = model(img)[0]  # only get inference result
        t2 = torch_utils.time_synchronized()
        print("inference time: {}s".format(t2 - t1))
        print('model: {}'.format(file))



        pred = utils.non_max_suppression(pred, conf_thres=0.1, iou_thres=0.6, multi_label=True)[0]
        t3 = time.time()
        print("post-processing time: {}s".format(t3 - t2))



        # process detections
        pred[:, :4] = utils.scale_coords(img.shape[2:], pred[:, :4], img_o.shape).round()


        bboxes = pred[:, :4].detach().cpu().numpy()
        scores = pred[:, 4].detach().cpu().numpy()
        classes = pred[:, 5].detach().cpu().numpy().astype(np.int) + 1

        pil_img = Image.fromarray(img_o[:, :, ::-1])
        plot_img = draw_objs(pil_img,
                             bboxes,
                             classes,
                             scores,
                             category_index=category_index,
                             box_thresh=0.2,
                             line_thickness=3,
                             font='arial.ttf',
                             font_size=30)
        plt.imshow(plot_img)
        plt.show()
        # 保存预测的图片结果

        plot_img.save(save_path)


if __name__ == "__main__":
    predict()
相关推荐
cver1233 小时前
野生动物检测数据集介绍-5,138张图片 野生动物保护监测 智能狩猎相机系统 生态研究与调查
人工智能·pytorch·深度学习·目标检测·计算机视觉·目标跟踪
学技术的大胜嗷3 小时前
离线迁移 Conda 环境到 Windows 服务器:用 conda-pack 摆脱硬路径限制
人工智能·深度学习·yolo·目标检测·机器学习
kyle~5 小时前
目标检测在国防和政府的应用实例
人工智能·目标检测·计算机视觉
weixin_377634848 小时前
【数据增强】精细化贴图数据增强
人工智能·目标检测·贴图
加油吧zkf9 小时前
目标检测新纪元:DETR到Mamba实战解析
图像处理·人工智能·python·目标检测·分类
一花·一叶1 天前
基于昇腾310B4的YOLOv8目标检测推理
yolo·目标检测·边缘计算
昵称是6硬币1 天前
YOLOv11: AN OVERVIEW OF THE KEY ARCHITECTURAL ENHANCEMENTS目标检测论文精读(逐段解析)
图像处理·人工智能·深度学习·yolo·目标检测·计算机视觉
OICQQ676580081 天前
创建一个基于YOLOv8+PyQt界面的驾驶员疲劳驾驶检测系统 实现对驾驶员疲劳状态的打哈欠检测,头部下垂 疲劳眼睛检测识别
yolo·pyqt·疲劳驾驶·检测识别·驾驶员检测·打哈欠检测·眼睛疲劳
云天徽上10 天前
【目标检测】图像处理基础:像素、分辨率与图像格式解析
图像处理·人工智能·目标检测·计算机视觉·数据可视化
king of code porter10 天前
目标检测之YOLOv5到YOLOv11——从架构设计和损失函数的变化分析
人工智能·yolo·目标检测