Yolov5 batch 推理

最新推荐文章于 2024-08-14 23:30:40 发布

犟小孩

最新推荐文章于 2024-08-14 23:30:40 发布

阅读量198

点赞数 1

分类专栏：代码查询文章标签： python 开发语言

本文链接：https://blog.csdn.net/qq_55542491/article/details/134264211

版权

前言

想要就有了

代码

import shutil
import time
import traceback
import torch

import os
import cv2
class PeopleDetect(object):
    def __init__(self, repo_or_dir, weight_path, confidence) -> None:
        self.model = torch.hub.load(repo_or_dir, "custom", path=weight_path, source='local',force_reload=True)
        self.confidence = confidence
    def start(self, img) -> list:
        mode_result = self.model(img,size=640)
        assert mode_result
        detect_result = mode_result.pandas().xyxy[0].to_dict(orient='index')
        if (len(detect_result)) > 0:
            box_list = self._get_box(detect_result)
            return box_list
        else:
            return []
    def _get_box(self, detect_result:list) -> list:
        box_list = []
        for i in range(len(detect_result)):
            conf = detect_result[i]['confidence']
            xm