modules/pytorch_jacinto_ai/xvision/datasets/pixel2pixel/coco_plus.py

   1 #################################################################################
   2 # Copyright (c) 2018-2021, Texas Instruments Incorporated - http://www.ti.com
   3 # All Rights Reserved.
   4 #
   5 # Redistribution and use in source and binary forms, with or without
   6 # modification, are permitted provided that the following conditions are met:
   7 #
   8 # * Redistributions of source code must retain the above copyright notice, this
   9 #   list of conditions and the following disclaimer.
  10 #
  11 # * Redistributions in binary form must reproduce the above copyright notice,
  12 #   this list of conditions and the following disclaimer in the documentation
  13 #   and/or other materials provided with the distribution.
  14 #
  15 # * Neither the name of the copyright holder nor the names of its
  16 #   contributors may be used to endorse or promote products derived from
  17 #   this software without specific prior written permission.
  18 #
  19 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  20 # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  21 # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  22 # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
  23 # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  24 # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
  25 # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
  26 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
  27 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  28 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  29 #
  30 #################################################################################
  31 # Also includes parts from: https://github.com/pytorch/vision
  32 # License: License: https://github.com/pytorch/vision/blob/master/LICENSE
  33
  34 # BSD 3-Clause License
  35 #
  36 # Copyright (c) Soumith Chintala 2016,
  37 # All rights reserved.
  38 #
  39 # Redistribution and use in source and binary forms, with or without
  40 # modification, are permitted provided that the following conditions are met:
  41 #
  42 # * Redistributions of source code must retain the above copyright notice, this
  43 #   list of conditions and the following disclaimer.
  44 #
  45 # * Redistributions in binary form must reproduce the above copyright notice,
  46 #   this list of conditions and the following disclaimer in the documentation
  47 #   and/or other materials provided with the distribution.
  48 #
  49 # * Neither the name of the copyright holder nor the names of its
  50 #   contributors may be used to endorse or promote products derived from
  51 #   this software without specific prior written permission.
  52 #
  53 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  54 # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  55 # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  56 # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
  57 # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  58 # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
  59 # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
  60 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
  61 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  62 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  63
  64 """
  65 Reference:
  66
  67 Microsoft COCO: Common Objects in Context,
  68 Tsung-Yi Lin, Michael Maire, Serge Belongie, Lubomir Bourdev, Ross Girshick, James Hays,
  69 Pietro Perona, Deva Ramanan, C. Lawrence Zitnick, Piotr Dollár,
  70 https://arxiv.org/abs/1405.0312, https://cocodataset.org/
  71 """
  72
  73 import numpy as np
  74 from .. import utils
  75 from ..coco import COCOSegmentation
  76 from .... import xnn
  77
  78 __all__ = ['coco_segmentation', 'coco_seg21']
  79
  80
  81 class COCOSegmentationPlus(COCOSegmentation):
  82     NUM_CLASSES = 80
  83     def __init__(self, *args, num_classes=NUM_CLASSES, transforms=None, **kwargs):
  84         # 21 class is a special case, otherwise use all the classes
  85         # in get_item a modulo is done to map the target to the required num_classes
  86         super().__init__(*args, num_classes=(num_classes if num_classes==21 else self.NUM_CLASSES), **kwargs)
  87         self.num_classes_ = num_classes
  88         self.void_classes = []
  89         self.valid_classes = range(self.num_classes_)
  90         self.ignore_index = 255
  91         self.class_map = dict(zip(self.valid_classes, range(self.num_classes_)))
  92         self.colors = utils.get_color_palette(num_classes)
  93         self.colors = (self.colors * self.num_classes_)[:self.num_classes_]
  94         self.label_colours = dict(zip(range(self.num_classes_), self.colors))
  95         self.transforms = transforms
  96
  97     def __getitem__(self, item):
  98         image, target = super().__getitem__(item)
  99         target = np.remainder(target, self.num_classes_)
 100         image = [image]
 101         target = [target]
 102         if self.transforms is not None:
 103             image, target = self.transforms(image, target)
 104         #
 105         return image, target
 106
 107     def num_classes(self):
 108         nc = []
 109         nc.append(self.num_classes_)
 110         return nc
 111
 112     def decode_segmap(self, temp):
 113         r = temp.copy()
 114         g = temp.copy()
 115         b = temp.copy()
 116         for l in range(0, self.num_classes_):
 117             r[temp == l] = self.label_colours[l][0]
 118             g[temp == l] = self.label_colours[l][1]
 119             b[temp == l] = self.label_colours[l][2]
 120         #
 121         rgb = np.zeros((temp.shape[0], temp.shape[1], 3))
 122         rgb[:, :, 0] = r / 255.0
 123         rgb[:, :, 1] = g / 255.0
 124         rgb[:, :, 2] = b / 255.0
 125         return rgb
 126
 127     def encode_segmap(self, mask):
 128         # Put all void classes to zero
 129         for _voidc in self.void_classes:
 130             mask[mask == _voidc] = self.ignore_index
 131         for _validc in self.valid_classes:
 132             mask[mask == _validc] = self.class_map[_validc]
 133         return mask
 134
 135
 136 ###########################################
 137 # config settings
 138 def get_config():
 139     dataset_config = xnn.utils.ConfigNode()
 140     dataset_config.num_classes = 80
 141     return dataset_config
 142
 143 def coco_segmentation(dataset_config, root, split=None, transforms=None, *args, **kwargs):
 144     dataset_config = get_config().merge_from(dataset_config)
 145     train_split = val_split = None
 146     split = ['train2017', 'val2017']
 147     for split_name in split:
 148         if split_name.startswith('train'):
 149             train_split = COCOSegmentationPlus(root, split_name, num_classes=dataset_config.num_classes,
 150                             transforms=transforms[0], *args, **kwargs)
 151         elif split_name.startswith('val'):
 152             val_split = COCOSegmentationPlus(root, split_name, num_classes=dataset_config.num_classes,
 153                             transforms=transforms[1], *args, **kwargs)
 154         else:
 155             assert False, 'unknown split'
 156         #
 157     #
 158     return train_split, val_split
 159
 160
 161 def coco_seg21(dataset_config, root, split=None, transforms=None, *args, **kwargs):
 162     dataset_config = get_config().merge_from(dataset_config)
 163     dataset_config.num_classes = 21
 164     return coco_segmentation(dataset_config, root, split=split, transforms=transforms, *args, **kwargs)
 165
 166