-
Notifications
You must be signed in to change notification settings - Fork 0
/
my_yolov6.py
133 lines (108 loc) · 5.68 KB
/
my_yolov6.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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
import torch
import numpy as np
import math
import cv2
from yolov6.utils.events import load_yaml
from yolov6.layers.common import DetectBackend
from yolov6.data.data_augment import letterbox
from yolov6.utils.nms import non_max_suppression
class my_yolov6():
def __init__(self, weights, device, yaml, img_size, half):
self.__dict__.update(locals())
# Init model
self.device = device
self.img_size = img_size
cuda = self.device != 'cpu' and torch.cuda.is_available()
self.device = torch.device(f'cuda:{device}' if cuda else 'cpu')
self.model = DetectBackend(weights, device=self.device)
self.stride = self.model.stride
self.class_names = load_yaml(yaml)['names']
self.img_size = self.check_img_size(self.img_size, s=self.stride) # check image size
# Half precision
if half & (self.device.type != 'cpu'):
self.model.model.half()
self.half = True
else:
self.model.model.float()
self.half = False
if self.device.type != 'cpu':
self.model(torch.zeros(1, 3, *self.img_size).to(self.device).type_as(
next(self.model.model.parameters()))) # warmup
# Switch model to deploy status
self.model_switch(self.model.model, self.img_size)
@staticmethod
def plot_box_and_label(image, lw, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255)):
# Add one xyxy box to image with label
p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
cv2.rectangle(image, p1, p2, color, thickness=lw, lineType=cv2.LINE_AA)
if label:
tf = max(lw - 1, 1) # font thickness
w, h = cv2.getTextSize(label, 0, fontScale=lw / 3, thickness=tf)[0] # text width, height
outside = p1[1] - h - 3 >= 0 # label fits outside box
p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
cv2.rectangle(image, p1, p2, color, -1, cv2.LINE_AA) # filled
cv2.putText(image, label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), 0, lw / 3, txt_color,
thickness=tf, lineType=cv2.LINE_AA)
@staticmethod
def rescale(ori_shape, boxes, target_shape):
'''Rescale the output to the original image shape'''
ratio = min(ori_shape[0] / target_shape[0], ori_shape[1] / target_shape[1])
padding = (ori_shape[1] - target_shape[1] * ratio) / 2, (ori_shape[0] - target_shape[0] * ratio) / 2
boxes[:, [0, 2]] -= padding[0]
boxes[:, [1, 3]] -= padding[1]
boxes[:, :4] /= ratio
boxes[:, 0].clamp_(0, target_shape[1]) # x1
boxes[:, 1].clamp_(0, target_shape[0]) # y1
boxes[:, 2].clamp_(0, target_shape[1]) # x2
boxes[:, 3].clamp_(0, target_shape[0]) # y2
return boxes
@staticmethod
def make_divisible(x, divisor):
# Upward revision the value x to make it evenly divisible by the divisor.
return math.ceil(x / divisor) * divisor
def check_img_size(self, img_size, s=32, floor=0):
"""Make sure image size is a multiple of stride s in each dimension, and return a new shape list of image."""
if isinstance(img_size, int): # integer i.e. img_size=640
new_size = max(self.make_divisible(img_size, int(s)), floor)
elif isinstance(img_size, list): # list i.e. img_size=[640, 480]
new_size = [max(self.make_divisible(x, int(s)), floor) for x in img_size]
else:
raise Exception(f"Unsupported type of img_size: {type(img_size)}")
if new_size != img_size:
print(f'WARNING: --img-size {img_size} must be multiple of max stride {s}, updating to {new_size}')
return new_size if isinstance(img_size,list) else [new_size]*2
def model_switch(self, model, img_size):
''' Model switch to deploy status '''
from yolov6.layers.common import RepVGGBlock
for layer in model.modules():
if isinstance(layer, RepVGGBlock):
layer.switch_to_deploy()
def precess_image(self,img_src, img_size, stride, half):
'''Process image before image inference.'''
image = letterbox(img_src, img_size, stride=stride)[0]
# Convert
image = image.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
image = torch.from_numpy(np.ascontiguousarray(image))
image = image.half() if half else image.float() # uint8 to fp16/32
image /= 255 # 0 - 255 to 0.0 - 1.0
return image, img_src
def infer(self, source, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic_nms=False, max_det=1000, interested_class=None):
img, img_src = self.precess_image(source, self.img_size, self.stride, self.half)
img = img.to(self.device)
if len(img.shape) == 3:
img = img[None]
# expand for batch dim
pred_results = self.model(img)
det = non_max_suppression(pred_results, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)[0]
class_names = []
if len(det):
det[:, :4] = self.rescale(img.shape[2:], det[:, :4], img_src.shape).round()
for *xyxy, conf, cls in reversed(det):
class_num = int(cls) # integer class
label = f'{self.class_names[class_num]} {conf:.2f}'
class_name = self.class_names[class_num]
if class_name in interested_class:
class_names.append(class_name)
self.plot_box_and_label(img_src, max(round(sum(img_src.shape) / 2 * 0.003), 2), xyxy, label, color=(255,0,0))
img_src = np.asarray(img_src)
return img_src, len(det), class_names