A SemanticSegmentor with test-time augmentation enabled. Its :meth:`__call__` method has the same interface as :meth:`SemanticSegmentor.forward`. combined with customized augmentation for original image
| 23 | |
| 24 | |
| 25 | class SemanticSegmentorWithTTA(nn.Module): |
| 26 | """ |
| 27 | A SemanticSegmentor with test-time augmentation enabled. |
| 28 | Its :meth:`__call__` method has the same interface as :meth:`SemanticSegmentor.forward`. |
| 29 | |
| 30 | combined with customized augmentation for original image |
| 31 | """ |
| 32 | |
| 33 | def __init__(self, cfg, model, batch_size=1): |
| 34 | """ |
| 35 | Args: |
| 36 | cfg (CfgNode): |
| 37 | model (SemanticSegmentor): a SemanticSegmentor to apply TTA on. |
| 38 | tta_mapper (callable): takes a dataset dict and returns a list of |
| 39 | augmented versions of the dataset dict. Defaults to |
| 40 | `DatasetMapperTTA(cfg)`. |
| 41 | batch_size (int): batch the augmented images into this batch size for inference. |
| 42 | """ |
| 43 | super().__init__() |
| 44 | if isinstance(model, DistributedDataParallel) or isinstance(model, DistModule): |
| 45 | model = model.module |
| 46 | self.cfg = cfg |
| 47 | |
| 48 | self.min_sizes = cfg.min_sizes |
| 49 | self.max_size = cfg.max_size |
| 50 | self.flip = cfg.flip |
| 51 | |
| 52 | self.model = model |
| 53 | |
| 54 | # if tta_mapper is None: |
| 55 | # tta_mapper = DatasetMapperTTA(cfg) |
| 56 | # self.tta_mapper = tta_mapper |
| 57 | assert batch_size == 1 |
| 58 | self.batch_size = batch_size |
| 59 | |
| 60 | def tta_mapper(self, dataset_dict): |
| 61 | """ |
| 62 | Args: |
| 63 | dict: a dict in standard model input format. See tutorials for details. |
| 64 | |
| 65 | Returns: |
| 66 | list[dict]: |
| 67 | a list of dicts, which contain augmented version of the input image. |
| 68 | The total number of dicts is ``len(min_sizes) * (2 if flip else 1)``. |
| 69 | Each dict has field "transforms" which is a TransformList, |
| 70 | containing the transforms that are used to generate this image. |
| 71 | """ |
| 72 | assert len(dataset_dict["image"].shape) == 4 |
| 73 | numpy_image = dataset_dict["image"].squeeze().permute(1, 2, 0).cpu().numpy() |
| 74 | shape = numpy_image.shape |
| 75 | orig_shape = (dataset_dict["height"], dataset_dict["width"]) |
| 76 | if shape[:2] != orig_shape: |
| 77 | # It transforms the "original" image in the dataset to the input image |
| 78 | pre_tfm = ResizeTransform(orig_shape[0], orig_shape[1], shape[0], shape[1]) |
| 79 | else: |
| 80 | pre_tfm = NoOpTransform() |
| 81 | |
| 82 | # Create all combinations of augmentations to use |
no outgoing calls