forked from airockchip/rknn_model_zoo
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathcoco_utils.py
176 lines (150 loc) · 7.06 KB
/
coco_utils.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
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
from copy import copy
import os
import cv2
import numpy as np
import json
class Letter_Box_Info():
def __init__(self, shape, new_shape, w_ratio, h_ratio, dw, dh, pad_color) -> None:
self.origin_shape = shape
self.new_shape = new_shape
self.w_ratio = w_ratio
self.h_ratio = h_ratio
self.dw = dw
self.dh = dh
self.pad_color = pad_color
def coco_eval_with_json(anno_json, pred_json):
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
anno = COCO(anno_json)
pred = anno.loadRes(pred_json)
eval = COCOeval(anno, pred, 'bbox')
# eval.params.useCats = 0
# eval.params.maxDets = list((100, 300, 1000))
# a = np.array(list(range(50, 96, 1)))/100
# eval.params.iouThrs = a
eval.evaluate()
eval.accumulate()
eval.summarize()
map, map50 = eval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5)
print('map --> ', map)
print('map50--> ', map50)
print('map75--> ', eval.stats[2])
print('map85--> ', eval.stats[-2])
print('map95--> ', eval.stats[-1])
class COCO_test_helper():
def __init__(self, enable_letter_box = False) -> None:
self.record_list = []
self.enable_ltter_box = enable_letter_box
if self.enable_ltter_box is True:
self.letter_box_info_list = []
else:
self.letter_box_info_list = None
def letter_box(self, im, new_shape, pad_color=(0,0,0), info_need=False):
# Resize and pad image while meeting stride-multiple constraints
shape = im.shape[:2] # current shape [height, width]
if isinstance(new_shape, int):
new_shape = (new_shape, new_shape)
# Scale ratio
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
# Compute padding
ratio = r # width, height ratios
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
dw /= 2 # divide padding into 2 sides
dh /= 2
if shape[::-1] != new_unpad: # resize
im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=pad_color) # add border
if self.enable_ltter_box is True:
self.letter_box_info_list.append(Letter_Box_Info(shape, new_shape, ratio, ratio, dw, dh, pad_color))
if info_need is True:
return im, ratio, (dw, dh)
else:
return im
def direct_resize(self, im, new_shape, info_need=False):
shape = im.shape[:2]
h_ratio = new_shape[0]/ shape[0]
w_ratio = new_shape[1]/ shape[1]
if self.enable_ltter_box is True:
self.letter_box_info_list.append(Letter_Box_Info(shape, new_shape, w_ratio, h_ratio, 0, 0, (0,0,0)))
im = cv2.resize(im, (new_shape[1], new_shape[0]))
return im
def get_real_box(self, box, in_format='xyxy'):
bbox = copy(box)
if self.enable_ltter_box == True:
# unletter_box result
if in_format=='xyxy':
bbox[:,0] -= self.letter_box_info_list[-1].dw
bbox[:,0] /= self.letter_box_info_list[-1].w_ratio
bbox[:,0] = np.clip(bbox[:,0], 0, self.letter_box_info_list[-1].origin_shape[1])
bbox[:,1] -= self.letter_box_info_list[-1].dh
bbox[:,1] /= self.letter_box_info_list[-1].h_ratio
bbox[:,1] = np.clip(bbox[:,1], 0, self.letter_box_info_list[-1].origin_shape[0])
bbox[:,2] -= self.letter_box_info_list[-1].dw
bbox[:,2] /= self.letter_box_info_list[-1].w_ratio
bbox[:,2] = np.clip(bbox[:,2], 0, self.letter_box_info_list[-1].origin_shape[1])
bbox[:,3] -= self.letter_box_info_list[-1].dh
bbox[:,3] /= self.letter_box_info_list[-1].h_ratio
bbox[:,3] = np.clip(bbox[:,3], 0, self.letter_box_info_list[-1].origin_shape[0])
return bbox
def get_real_seg(self, seg):
#! fix side effect
dh = int(self.letter_box_info_list[-1].dh)
dw = int(self.letter_box_info_list[-1].dw)
origin_shape = self.letter_box_info_list[-1].origin_shape
new_shape = self.letter_box_info_list[-1].new_shape
if (dh == 0) and (dw == 0) and origin_shape == new_shape:
return seg
elif dh == 0 and dw != 0:
seg = seg[:, :, dw:-dw] # a[0:-0] = []
elif dw == 0 and dh != 0 :
seg = seg[:, dh:-dh, :]
seg = np.where(seg, 1, 0).astype(np.uint8).transpose(1,2,0)
seg = cv2.resize(seg, (origin_shape[1], origin_shape[0]), interpolation=cv2.INTER_LINEAR)
if len(seg.shape) < 3:
return seg[None,:,:]
else:
return seg.transpose(2,0,1)
def add_single_record(self, image_id, category_id, bbox, score, in_format='xyxy', pred_masks = None):
if self.enable_ltter_box == True:
# unletter_box result
if in_format=='xyxy':
bbox[0] -= self.letter_box_info_list[-1].dw
bbox[0] /= self.letter_box_info_list[-1].w_ratio
bbox[1] -= self.letter_box_info_list[-1].dh
bbox[1] /= self.letter_box_info_list[-1].h_ratio
bbox[2] -= self.letter_box_info_list[-1].dw
bbox[2] /= self.letter_box_info_list[-1].w_ratio
bbox[3] -= self.letter_box_info_list[-1].dh
bbox[3] /= self.letter_box_info_list[-1].h_ratio
# bbox = [value/self.letter_box_info_list[-1].ratio for value in bbox]
if in_format=='xyxy':
# change xyxy to xywh
bbox[2] = bbox[2] - bbox[0]
bbox[3] = bbox[3] - bbox[1]
else:
assert False, "now only support xyxy format, please add code to support others format"
def single_encode(x):
from pycocotools.mask import encode
rle = encode(np.asarray(x[:, :, None], order="F", dtype="uint8"))[0]
rle["counts"] = rle["counts"].decode("utf-8")
return rle
if pred_masks is None:
self.record_list.append({"image_id": image_id,
"category_id": category_id,
"bbox":[round(x, 3) for x in bbox],
'score': round(score, 5),
})
else:
rles = single_encode(pred_masks)
self.record_list.append({"image_id": image_id,
"category_id": category_id,
"bbox":[round(x, 3) for x in bbox],
'score': round(score, 5),
'segmentation': rles,
})
def export_to_json(self, path):
with open(path, 'w') as f:
json.dump(self.record_list, f)