This repository has been archived by the owner on Jan 26, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 1
/
my_detect.py
82 lines (67 loc) · 2.84 KB
/
my_detect.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
import random
import numpy as np
from yolov5.utils.datasets import letterbox
import torch
from torch.backends import cudnn
from yolov5.models.experimental import attempt_load
from yolov5.utils.general import check_img_size, non_max_suppression, scale_coords, xyxy2xywh
from yolov5.utils.torch_utils import select_device
import sys
# sys.path.insert(0, './yolov5')
class _DetectOpt:
def __init__(self):
self.output = 'inference/output'
self.source = 0
# self.deep_sort_weights = 'deep_sort_pytorch/deep_sort/deep/checkpoint/ckpt.t7'
self.yolo_weights = 'yolov5/weights/yolov5s.pt'
self.imgsz = 640
self.evaluate = False
self.device = '0'
# self.config_deepsort = "deep_sort_pytorch/configs/deep_sort.yaml"
self.conf_thres = 0.5
self.iou_thres = 0.7
self.classes = 0
self.agnostic_nms = False
class MyDetect:
def __init__(self):
self.opt = _DetectOpt()
self._load_model()
def _load_model(self):
self.device = select_device(self.opt.device)
self.half = self.device.type != 'cpu' # half precision only supported on CUDA
self.model = attempt_load(self.opt.yolo_weights, map_location=self.device) # load FP32 model
self.stride = int(self.model.stride.max()) # model stride
self.imgsz = check_img_size(self.opt.imgsz, s=self.stride) # check img_size
if self.half:
self.model.half()
cudnn.benchmark = True # set True to speed up constant image size inference
self.names = self.model.module.names if hasattr(self.model, 'module') else self.model.names
self.colors = [[random.randint(0, 255) for _ in range(3)] for _ in self.names]
pass
def _process_image(self, source):
pro_image = letterbox(source, new_shape=self.opt.imgsz)[0]
pro_image = pro_image[:, :, ::-1].transpose(2, 0, 1)
pro_image = np.ascontiguousarray(pro_image)
pro_image = torch.from_numpy(pro_image).to(self.device)
pro_image = pro_image.half()
pro_image /= 255
if pro_image.ndimension() == 3:
pro_image = pro_image.unsqueeze(0)
return pro_image
def detect(self, source):
"""
返回xywhs,confs,cls_s 这些都是tensor
"""
xywhs = confs = clss = None
processed_image = self._process_image(source)
pred = self.model(processed_image, augment=False)[0]
pred = non_max_suppression(
pred, self.opt.conf_thres, self.opt.iou_thres)
for det in pred:
if det is not None and len(det):
det[:, :4] = scale_coords(
processed_image.shape[2:], det[:, :4], source.shape).round()
xywhs = xyxy2xywh(det[:, 0:4])
confs = det[:, 4]
clss = det[:, 5]
return xywhs, confs, clss