Views
No views yet
| Model | Download | Download (with sample test data) | ONNX version | Opset version | Accuracy |
|---|---|---|---|---|---|
| YOLOv4 | 251 MB | 236 MB | 1.6 | 11 | mAP of 0.5733 |
(1, 416, 416, 3). Each dimension represents the following variables: (batch_size, height, width, channels).1import numpy as np
2import cv2
3
4# this function is from tensorflow-yolov4-tflite/core/utils.py
5def image_preprocess(image, target_size, gt_boxes=None):
6
7ih, iw = target_size
8h, w, _ = image.shape
9
10scale = min(iw/w, ih/h)
11nw, nh = int(scale * w), int(scale * h)
12image_resized = cv2.resize(image, (nw, nh))
13
14image_padded = np.full(shape=[ih, iw, 3], fill_value=128.0)
15dw, dh = (iw - nw) // 2, (ih-nh) // 2
16image_padded[dh:nh+dh, dw:nw+dw, :] = image_resized
17image_padded = image_padded / 255.
18
19if gt_boxes is None:
20return image_padded
21
22else:
23gt_boxes[:, [0, 2]] = gt_boxes[:, [0, 2]] * scale + dw
24gt_boxes[:, [1, 3]] = gt_boxes[:, [1, 3]] * scale + dh
25return image_padded, gt_boxes
26
27# input
28input_size = 416
29
30original_image = cv2.imread("input.jpg")
31original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
32original_image_size = original_image.shape[:2]
33
34image_data = image_preprocess(np.copy(original_image), [input_size, input_size])
35image_data = image_data[np.newaxis, ...].astype(np.float32)
36(1, 52, 52, 3, 85)1from scipy import special
2import colorsys
3import random
4
5
6def get_anchors(anchors_path, tiny=False):
7'''loads the anchors from a file'''
8with open(anchors_path) as f:
9anchors = f.readline()
10anchors = np.array(anchors.split(','), dtype=np.float32)
11return anchors.reshape(3, 3, 2)
12
13def postprocess_bbbox(pred_bbox, ANCHORS, STRIDES, XYSCALE=[1,1,1]):
14'''define anchor boxes'''
15for i, pred in enumerate(pred_bbox):
16conv_shape = pred.shape
17output_size = conv_shape[1]
18conv_raw_dxdy = pred[:, :, :, :, 0:2]
19conv_raw_dwdh = pred[:, :, :, :, 2:4]
20xy_grid = np.meshgrid(np.arange(output_size), np.arange(output_size))
21xy_grid = np.expand_dims(np.stack(xy_grid, axis=-1), axis=2)
22
23xy_grid = np.tile(np.expand_dims(xy_grid, axis=0), [1, 1, 1, 3, 1])
24xy_grid = xy_grid.astype(np.float)
25
26pred_xy = ((special.expit(conv_raw_dxdy) * XYSCALE[i]) - 0.5 * (XYSCALE[i] - 1) + xy_grid) * STRIDES[i]
27pred_wh = (np.exp(conv_raw_dwdh) * ANCHORS[i])
28pred[:, :, :, :, 0:4] = np.concatenate([pred_xy, pred_wh], axis=-1)
29
30pred_bbox = [np.reshape(x, (-1, np.shape(x)[-1])) for x in pred_bbox]
31pred_bbox = np.concatenate(pred_bbox, axis=0)
32return pred_bbox
33
34
35def postprocess_boxes(pred_bbox, org_img_shape, input_size, score_threshold):
36'''remove boundary boxs with a low detection probability'''
37valid_scale=[0, np.inf]
38pred_bbox = np.array(pred_bbox)
39
40pred_xywh = pred_bbox[:, 0:4]
41pred_conf = pred_bbox[:, 4]
42pred_prob = pred_bbox[:, 5:]
43
44# (1) (x, y, w, h) --> (xmin, ymin, xmax, ymax)
45pred_coor = np.concatenate([pred_xywh[:, :2] - pred_xywh[:, 2:] * 0.5,
46pred_xywh[:, :2] + pred_xywh[:, 2:] * 0.5], axis=-1)
47# (2) (xmin, ymin, xmax, ymax) -> (xmin_org, ymin_org, xmax_org, ymax_org)
48org_h, org_w = org_img_shape
49resize_ratio = min(input_size / org_w, input_size / org_h)
50
51dw = (input_size - resize_ratio * org_w) / 2
52dh = (input_size - resize_ratio * org_h) / 2
53
54pred_coor[:, 0::2] = 1.0 * (pred_coor[:, 0::2] - dw) / resize_ratio
55pred_coor[:, 1::2] = 1.0 * (pred_coor[:, 1::2] - dh) / resize_ratio
56
57# (3) clip some boxes that are out of range
58pred_coor = np.concatenate([np.maximum(pred_coor[:, :2], [0, 0]),
59np.minimum(pred_coor[:, 2:], [org_w - 1, org_h - 1])], axis=-1)
60invalid_mask = np.logical_or((pred_coor[:, 0] > pred_coor[:, 2]), (pred_coor[:, 1] > pred_coor[:, 3]))
61pred_coor[invalid_mask] = 0
62
63# (4) discard some invalid boxes
64bboxes_scale = np.sqrt(np.multiply.reduce(pred_coor[:, 2:4] - pred_coor[:, 0:2], axis=-1))
65scale_mask = np.logical_and((valid_scale[0] < bboxes_scale), (bboxes_scale < valid_scale[1]))
66
67# (5) discard some boxes with low scores
68classes = np.argmax(pred_prob, axis=-1)
69scores = pred_conf * pred_prob[np.arange(len(pred_coor)), classes]
70score_mask = scores > score_threshold
71mask = np.logical_and(scale_mask, score_mask)
72coors, scores, classes = pred_coor[mask], scores[mask], classes[mask]
73
74return np.concatenate([coors, scores[:, np.newaxis], classes[:, np.newaxis]], axis=-1)
75
76def bboxes_iou(boxes1, boxes2):
77'''calculate the Intersection Over Union value'''
78boxes1 = np.array(boxes1)
79boxes2 = np.array(boxes2)
80
81boxes1_area = (boxes1[..., 2] - boxes1[..., 0]) * (boxes1[..., 3] - boxes1[..., 1])
82boxes2_area = (boxes2[..., 2] - boxes2[..., 0]) * (boxes2[..., 3] - boxes2[..., 1])
83
84left_up = np.maximum(boxes1[..., :2], boxes2[..., :2])
85right_down = np.minimum(boxes1[..., 2:], boxes2[..., 2:])
86
87inter_section = np.maximum(right_down - left_up, 0.0)
88inter_area = inter_section[..., 0] * inter_section[..., 1]
89union_area = boxes1_area + boxes2_area - inter_area
90ious = np.maximum(1.0 * inter_area / union_area, np.finfo(np.float32).eps)
91
92return ious
93
94def nms(bboxes, iou_threshold, sigma=0.3, method='nms'):
95"""
96:param bboxes: (xmin, ymin, xmax, ymax, score, class)
97
98Note: soft-nms, https://arxiv.org/pdf/1704.04503.pdf
99https://github.com/bharatsingh430/soft-nms
100"""
101classes_in_img = list(set(bboxes[:, 5]))
102best_bboxes = []
103
104for cls in classes_in_img:
105cls_mask = (bboxes[:, 5] == cls)
106cls_bboxes = bboxes[cls_mask]
107
108while len(cls_bboxes) > 0:
109max_ind = np.argmax(cls_bboxes[:, 4])
110best_bbox = cls_bboxes[max_ind]
111best_bboxes.append(best_bbox)
112cls_bboxes = np.concatenate([cls_bboxes[: max_ind], cls_bboxes[max_ind + 1:]])
113iou = bboxes_iou(best_bbox[np.newaxis, :4], cls_bboxes[:, :4])
114weight = np.ones((len(iou),), dtype=np.float32)
115
116assert method in ['nms', 'soft-nms']
117
118if method == 'nms':
119iou_mask = iou > iou_threshold
120weight[iou_mask] = 0.0
121
122if method == 'soft-nms':
123weight = np.exp(-(1.0 * iou ** 2 / sigma))
124
125cls_bboxes[:, 4] = cls_bboxes[:, 4] * weight
126score_mask = cls_bboxes[:, 4] > 0.
127cls_bboxes = cls_bboxes[score_mask]
128
129return best_bboxes
130
131def read_class_names(class_file_name):
132'''loads class name from a file'''
133names = {}
134with open(class_file_name, 'r') as data:
135for ID, name in enumerate(data):
136names[ID] = name.strip('\n')
137return names
138
139def draw_bbox(image, bboxes, classes=read_class_names("coco.names"), show_label=True):
140"""
141bboxes: [x_min, y_min, x_max, y_max, probability, cls_id] format coordinates.
142"""
143
144num_classes = len(classes)
145image_h, image_w, _ = image.shape
146hsv_tuples = [(1.0 * x / num_classes, 1., 1.) for x in range(num_classes)]
147colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
148colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors))
149
150random.seed(0)
151random.shuffle(colors)
152random.seed(None)
153
154for i, bbox in enumerate(bboxes):
155coor = np.array(bbox[:4], dtype=np.int32)
156fontScale = 0.5
157score = bbox[4]
158class_ind = int(bbox[5])
159bbox_color = colors[class_ind]
160bbox_thick = int(0.6 * (image_h + image_w) / 600)
161c1, c2 = (coor[0], coor[1]), (coor[2], coor[3])
162cv2.rectangle(image, c1, c2, bbox_color, bbox_thick)
163
164if show_label:
165bbox_mess = '%s: %.2f' % (classes[class_ind], score)
166t_size = cv2.getTextSize(bbox_mess, 0, fontScale, thickness=bbox_thick//2)[0]
167cv2.rectangle(image, c1, (c1[0] + t_size[0], c1[1] - t_size[1] - 3), bbox_color, -1)
168
169cv2.putText(image, bbox_mess, (c1[0], c1[1]-2), cv2.FONT_HERSHEY_SIMPLEX,
170fontScale, (0, 0, 0), bbox_thick//2, lineType=cv2.LINE_AA)
171
172return image