tools/accuracy_checker/accuracy_checker/postprocessor/correct_yolo_v2_boxes.py

   1 """
   2 Copyright (c) 2019 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
  17 from ..config import NumberField
  18 from .postprocessor import BasePostprocessorConfig, Postprocessor
  19 from ..representation import DetectionPrediction, DetectionAnnotation
  20 from ..utils import get_size_from_config
  21
  22
  23 class CorrectYoloV2Boxes(Postprocessor):
  24     __provider__ = 'correct_yolo_v2_boxes'
  25
  26     prediction_types = (DetectionPrediction, )
  27     annotation_types = (DetectionAnnotation, )
  28
  29     def validate_config(self):
  30         class _CorrectYoloV2BoxesConfigValidator(BasePostprocessorConfig):
  31             dst_width = NumberField(floats=False, optional=True, min_value=1)
  32             dst_height = NumberField(floats=False, optional=True, min_value=1)
  33             size = NumberField(floats=False, optional=True, min_value=1)
  34
  35         clip_config_validator = _CorrectYoloV2BoxesConfigValidator(
  36             self.__provider__, on_extra_argument=_CorrectYoloV2BoxesConfigValidator.ERROR_ON_EXTRA_ARGUMENT
  37         )
  38         clip_config_validator.validate(self.config)
  39
  40     def configure(self):
  41         self.dst_height, self.dst_width = get_size_from_config(self.config)
  42
  43     def process_image(self, annotation, prediction):
  44         dst_h, dst_w = self.dst_height, self.dst_width
  45         # postprocessor always expects lists of annotations and predictions for the same image
  46         # we do not need to get image sizes in cycle, because they are equal
  47         img_h, img_w, _ = self.image_size
  48
  49         if (dst_w / img_w) < (dst_h / img_h):
  50             new_w = dst_w
  51             new_h = (img_h * dst_w) // img_w
  52         else:
  53             new_h = dst_h
  54             new_w = (img_w * dst_h) // img_h
  55
  56         for prediction_ in prediction:
  57             coordinates = zip(prediction_.x_mins, prediction_.y_mins, prediction_.x_maxs, prediction_.y_maxs)
  58             for i, (x0, y0, x1, y1) in enumerate(coordinates):
  59                 box = [(x0 + x1) / 2.0, (y0 + y1) / 2.0, x1 - x0, y1 - y0]
  60                 box[0] = (box[0] - (dst_w - new_w) / (2.0 * dst_w)) * (dst_w / new_w)
  61                 box[1] = (box[1] - (dst_h - new_h) / (2.0 * dst_h)) * (dst_h / new_h)
  62                 box[2] *= dst_w / new_w
  63                 box[3] *= dst_h / new_h
  64
  65                 box[0] *= img_w
  66                 box[1] *= img_h
  67                 box[2] *= img_w
  68                 box[3] *= img_h
  69
  70                 prediction_.x_mins[i] = box[0] - box[2] / 2.0 + 1
  71                 prediction_.y_mins[i] = box[1] - box[3] / 2.0 + 1
  72                 prediction_.x_maxs[i] = box[0] + box[2] / 2.0 + 1
  73                 prediction_.y_maxs[i] = box[1] + box[3] / 2.0 + 1
  74
  75         return annotation, prediction