| 176 | |
| 177 | |
| 178 | class NanoDetABC(metaclass=ABCMeta): |
| 179 | def __init__( |
| 180 | self, |
| 181 | input_shape=[272, 160], |
| 182 | reg_max=7, |
| 183 | strides=[8, 16, 32], |
| 184 | prob_threshold=0.4, |
| 185 | iou_threshold=0.3, |
| 186 | num_candidate=1000, |
| 187 | top_k=-1, |
| 188 | class_names=["face"], |
| 189 | ): |
| 190 | self.strides = strides |
| 191 | self.input_shape = input_shape |
| 192 | self.reg_max = reg_max |
| 193 | self.prob_threshold = prob_threshold |
| 194 | self.iou_threshold = iou_threshold |
| 195 | self.num_candidate = num_candidate |
| 196 | self.top_k = top_k |
| 197 | self.img_mean = [103.53, 116.28, 123.675] |
| 198 | self.img_std = [57.375, 57.12, 58.395] |
| 199 | self.input_size = (self.input_shape[1], self.input_shape[0]) |
| 200 | self.class_names = class_names |
| 201 | self.num_classes = len(self.class_names) |
| 202 | |
| 203 | def preprocess(self, img): |
| 204 | # resize image |
| 205 | ResizeM = get_resize_matrix((img.shape[1], img.shape[0]), self.input_size, True) |
| 206 | img_resize = cv2.warpPerspective(img, ResizeM, dsize=self.input_size) |
| 207 | |
| 208 | # normalize image |
| 209 | img_input = img_resize.astype(np.float32) / 255 |
| 210 | img_mean = np.array(self.img_mean, dtype=np.float32).reshape(1, 1, 3) / 255 |
| 211 | img_std = np.array(self.img_std, dtype=np.float32).reshape(1, 1, 3) / 255 |
| 212 | img_input = (img_input - img_mean) / img_std |
| 213 | |
| 214 | # expand dims |
| 215 | img_input = np.transpose(img_input, [2, 0, 1]) |
| 216 | img_input = np.expand_dims(img_input, axis=0) |
| 217 | return img_input, ResizeM |
| 218 | |
| 219 | def postprocess(self, scores, raw_boxes, ResizeM, raw_shape): |
| 220 | # generate centers |
| 221 | decode_boxes = [] |
| 222 | select_scores = [] |
| 223 | for stride, box_distribute, score in zip(self.strides, raw_boxes, scores): |
| 224 | # centers |
| 225 | fm_h = self.input_shape[0] / stride |
| 226 | fm_w = self.input_shape[1] / stride |
| 227 | |
| 228 | h_range = np.arange(fm_h) |
| 229 | w_range = np.arange(fm_w) |
| 230 | ww, hh = np.meshgrid(w_range, h_range) |
| 231 | |
| 232 | ct_row = hh.flatten() * stride |
| 233 | ct_col = ww.flatten() * stride |
| 234 | |
| 235 | center = np.stack((ct_col, ct_row, ct_col, ct_row), axis=1) |
nothing calls this directly
no outgoing calls
no test coverage detected