前言
想要就有了
代码
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