|
| 1 | +""" |
| 2 | + Copyright (c) 2022 Intel Corporation |
| 3 | +
|
| 4 | + Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | + you may not use this file except in compliance with the License. |
| 6 | + You may obtain a copy of the License at |
| 7 | +
|
| 8 | + http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +
|
| 10 | + Unless required by applicable law or agreed to in writing, software |
| 11 | + distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | + See the License for the specific language governing permissions and |
| 14 | + limitations under the License. |
| 15 | +""" |
| 16 | +import math |
| 17 | +import numpy as np |
| 18 | + |
| 19 | +from .types import NumericalValue |
| 20 | +from .detection_model import DetectionModel |
| 21 | +from .utils import Detection, softmax, nms, clip_detections |
| 22 | + |
| 23 | + |
| 24 | +class NanoDet(DetectionModel): |
| 25 | + __model__ = 'NanoDet' |
| 26 | + |
| 27 | + def __init__(self, model_adapter, configuration=None, preload=False): |
| 28 | + super().__init__(model_adapter, configuration, preload) |
| 29 | + self._check_io_number(1, 1) |
| 30 | + self.output_blob_name = self._get_outputs() |
| 31 | + self.reg_max = 7 |
| 32 | + self.strides = [8, 16, 32] |
| 33 | + self.ad = 0.5 |
| 34 | + |
| 35 | + def _get_outputs(self): |
| 36 | + output_blob_name = next(iter(self.outputs)) |
| 37 | + output_size = self.outputs[output_blob_name].shape |
| 38 | + if len(output_size) != 3: |
| 39 | + self.raise_error("Unexpected output blob shape {}. Only 3D output blob is supported".format(output_size)) |
| 40 | + |
| 41 | + return output_blob_name |
| 42 | + |
| 43 | + @classmethod |
| 44 | + def parameters(cls): |
| 45 | + parameters = super().parameters() |
| 46 | + parameters['resize_type'].update_default_value('fit_to_window') |
| 47 | + parameters['confidence_threshold'].update_default_value(0.5) |
| 48 | + parameters.update({ |
| 49 | + 'iou_threshold': NumericalValue(default_value=0.6, description="Threshold for NMS filtering"), |
| 50 | + 'num_classes': NumericalValue(default_value=80, value_type=int, description="Number of classes") |
| 51 | + }) |
| 52 | + return parameters |
| 53 | + |
| 54 | + def postprocess(self, outputs, meta): |
| 55 | + detections = self._parse_outputs(outputs, meta) |
| 56 | + detections = self.rescale_detections(detections, meta) |
| 57 | + return detections |
| 58 | + |
| 59 | + def _parse_outputs(self, outputs, meta): |
| 60 | + output = outputs[self.output_blob_name][0] |
| 61 | + |
| 62 | + cls_scores = output[:, :self.num_classes] |
| 63 | + bbox_preds = output[:, self.num_classes:] |
| 64 | + input_height, input_width = meta['padded_shape'][:2] if meta.get('padded_shape') else meta['resized_shape'][:2] |
| 65 | + |
| 66 | + bboxes = self.get_bboxes(bbox_preds, input_height, input_width) |
| 67 | + dets = [] |
| 68 | + for label, score in enumerate(np.transpose(cls_scores)): |
| 69 | + mask = score > self.confidence_threshold |
| 70 | + filtered_boxes, score = bboxes[mask, :], score[mask] |
| 71 | + if score.size == 0: |
| 72 | + continue |
| 73 | + x_mins, y_mins, x_maxs, y_maxs = filtered_boxes.T |
| 74 | + keep = nms(x_mins, y_mins, x_maxs, y_maxs, score, self.iou_threshold, include_boundaries=True) |
| 75 | + score = score[keep] |
| 76 | + x_mins, y_mins, x_maxs, y_maxs = x_mins[keep], y_mins[keep], x_maxs[keep], y_maxs[keep] |
| 77 | + labels = np.full_like(score, label, dtype=int) |
| 78 | + dets += [Detection(*det) for det in zip(x_mins, y_mins, x_maxs, y_maxs, score, labels)] |
| 79 | + return dets |
| 80 | + |
| 81 | + @staticmethod |
| 82 | + def distance2bbox(points, distance, max_shape): |
| 83 | + x1 = np.expand_dims(points[:, 0] - distance[:, 0], -1).clip(0, max_shape[1]) |
| 84 | + y1 = np.expand_dims(points[:, 1] - distance[:, 1], -1).clip(0, max_shape[0]) |
| 85 | + x2 = np.expand_dims(points[:, 0] + distance[:, 2], -1).clip(0, max_shape[1]) |
| 86 | + y2 = np.expand_dims(points[:, 1] + distance[:, 3], -1).clip(0, max_shape[0]) |
| 87 | + return np.concatenate((x1, y1, x2, y2), axis=-1) |
| 88 | + |
| 89 | + def get_single_level_center_point(self, featmap_size, stride): |
| 90 | + h, w = featmap_size |
| 91 | + x_range, y_range = (np.arange(w) + self.ad) * stride, (np.arange(h) + self.ad) * stride |
| 92 | + y, x = np.meshgrid(y_range, x_range, indexing='ij') |
| 93 | + return y.flatten(), x.flatten() |
| 94 | + |
| 95 | + def get_bboxes(self, reg_preds, input_height, input_width): |
| 96 | + featmap_sizes = [(math.ceil(input_height / stride), math.ceil(input_width) / stride) for stride in self.strides] |
| 97 | + list_center_priors = [] |
| 98 | + for stride, featmap_size in zip(self.strides, featmap_sizes): |
| 99 | + y, x = self.get_single_level_center_point(featmap_size, stride) |
| 100 | + strides = np.full_like(x, stride) |
| 101 | + list_center_priors.append(np.stack([x, y, strides, strides], axis=-1)) |
| 102 | + center_priors = np.concatenate(list_center_priors, axis=0) |
| 103 | + dist_project = np.linspace(0, self.reg_max, self.reg_max + 1) |
| 104 | + x = np.dot(softmax(np.reshape(reg_preds, (*reg_preds.shape[:-1], 4, self.reg_max + 1)), -1, True), dist_project) |
| 105 | + dis_preds = x * np.expand_dims(center_priors[:, 2], -1) |
| 106 | + return self.distance2bbox(center_priors[:, :2], dis_preds, (input_height, input_width)) |
| 107 | + |
| 108 | + @staticmethod |
| 109 | + def rescale_detections(detections, meta): |
| 110 | + input_h, input_w, _ = meta['resized_shape'] |
| 111 | + orig_h, orig_w, _ = meta['original_shape'] |
| 112 | + w = orig_w / input_w |
| 113 | + h = orig_h / input_h |
| 114 | + |
| 115 | + for detection in detections: |
| 116 | + detection.xmin *= w |
| 117 | + detection.xmax *= w |
| 118 | + detection.ymin *= h |
| 119 | + detection.ymax *= h |
| 120 | + |
| 121 | + return clip_detections(detections, meta['original_shape']) |
| 122 | + |
| 123 | + |
| 124 | +class NanoDetPlus(NanoDet): |
| 125 | + __model__ = 'NanoDet-Plus' |
| 126 | + |
| 127 | + def __init__(self, model_adapter, configuration=None, preload=False): |
| 128 | + super().__init__(model_adapter, configuration, preload) |
| 129 | + self.ad = 0 |
| 130 | + self.strides = [8, 16, 32, 64] |
0 commit comments