怎么把yolov5的图片检测封装成类
时间: 2023-06-26 15:09:53 浏览: 77
可以创建一个 `YoloV5Detector` 类来封装 YOLOv5 的图片检测功能。下面是一个基本的示例代码:
```python
import torch
import cv2
import numpy as np
from models.experimental import attempt_load
from utils.general import non_max_suppression
class YoloV5Detector:
def __init__(self, model_path, conf_thresh=0.25, iou_thresh=0.45):
self.conf_thresh = conf_thresh
self.iou_thresh = iou_thresh
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = attempt_load(model_path, map_location=self.device)
self.model.eval()
def detect(self, image_path):
img = cv2.imread(image_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = img.transpose(2, 0, 1)
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(self.device).float() / 255.0
# Run inference
with torch.no_grad():
results = self.model(img, size=img.shape[-2:])
results = non_max_suppression(results, conf_thres=self.conf_thresh, iou_thres=self.iou_thresh)
return results
```
在上面的代码中,`YoloV5Detector` 类包含一个初始化方法,该方法加载预训练模型,并设置一些检测阈值。`detect` 方法接受图像路径作为输入,并返回一个字典列表,每个字典包含检测到的边界框的坐标、置信度和类别。
要使用 `YoloV5Detector` 类,只需实例化该类并调用 `detect` 方法即可:
```python
detector = YoloV5Detector('/path/to/yolov5s.pt')
results = detector.detect('/path/to/image.jpg')
```