# Copyright (c) Facebook, Inc. and its affiliates. # Copyright (c) Meta Platforms, Inc. All Rights Reserved import copy from itertools import count import math import numpy as np import torch from fvcore.transforms import HFlipTransform from torch import nn from torch.nn.parallel import DistributedDataParallel from detectron2.data.detection_utils import read_image from detectron2.modeling import DatasetMapperTTA from detectron2.modeling.postprocessing import sem_seg_postprocess import logging from detectron2.utils.logger import log_every_n, log_first_n __all__ = [ "SemanticSegmentorWithTTA", ] class SemanticSegmentorWithTTA(nn.Module): """ A SemanticSegmentor with test-time augmentation enabled. Its :meth:`__call__` method has the same interface as :meth:`SemanticSegmentor.forward`. """ def __init__(self, cfg, model, tta_mapper=None, batch_size=1): """ Args: cfg (CfgNode): model (SemanticSegmentor): a SemanticSegmentor to apply TTA on. tta_mapper (callable): takes a dataset dict and returns a list of augmented versions of the dataset dict. Defaults to `DatasetMapperTTA(cfg)`. batch_size (int): batch the augmented images into this batch size for inference. """ super().__init__() if isinstance(model, DistributedDataParallel): model = model.module self.cfg = cfg.clone() self.model = model if tta_mapper is None: tta_mapper = DatasetMapperTTA(cfg) self.tta_mapper = tta_mapper self.batch_size = batch_size def _inference_with_model(self, inputs): if self.cfg.TEST.SLIDING_WINDOW: log_first_n(logging.INFO, "Using sliding window to test") outputs = [] for input in inputs: image_size = input["image"].shape[1:] # h,w if self.cfg.TEST.SLIDING_TILE_SIZE > 0: tile_size = ( self.cfg.TEST.SLIDING_TILE_SIZE, self.cfg.TEST.SLIDING_TILE_SIZE, ) else: selected_mapping = {256: 224, 512: 256, 768: 512, 896: 512} tile_size = min(image_size) tile_size = selected_mapping[tile_size] tile_size = (tile_size, tile_size) extra_info = { k: v for k, v in input.items() if k not in ["image", "height", "width"] } log_every_n( logging.INFO, "split {} to {}".format(image_size, tile_size) ) overlap = self.cfg.TEST.SLIDING_OVERLAP stride = math.ceil(tile_size[0] * (1 - overlap)) tile_rows = int( math.ceil((image_size[0] - tile_size[0]) / stride) + 1 ) # strided convolution formula tile_cols = int(math.ceil((image_size[1] - tile_size[1]) / stride) + 1) full_probs = None count_predictions = None tile_counter = 0 for row in range(tile_rows): for col in range(tile_cols): x1 = int(col * stride) y1 = int(row * stride) x2 = min(x1 + tile_size[1], image_size[1]) y2 = min(y1 + tile_size[0], image_size[0]) x1 = max( int(x2 - tile_size[1]), 0 ) # for portrait images the x1 underflows sometimes y1 = max( int(y2 - tile_size[0]), 0 ) # for very few rows y1 underflows img = input["image"][:, y1:y2, x1:x2] padded_img = nn.functional.pad( img, ( 0, tile_size[1] - img.shape[-1], 0, tile_size[0] - img.shape[-2], ), ) tile_counter += 1 padded_input = {"image": padded_img} padded_input.update(extra_info) padded_prediction = self.model([padded_input])[0]["sem_seg"] prediction = padded_prediction[ :, 0 : img.shape[1], 0 : img.shape[2] ] if full_probs is None: full_probs = prediction.new_zeros( prediction.shape[0], image_size[0], image_size[1] ) if count_predictions is None: count_predictions = prediction.new_zeros( prediction.shape[0], image_size[0], image_size[1] ) count_predictions[:, y1:y2, x1:x2] += 1 full_probs[ :, y1:y2, x1:x2 ] += prediction # accumulate the predictions also in the overlapping regions full_probs /= count_predictions full_probs = sem_seg_postprocess( full_probs, image_size, input.get("height", image_size[0]), input.get("width", image_size[1]), ) outputs.append({"sem_seg": full_probs}) return outputs else: log_first_n(logging.INFO, "Using whole image to test") return self.model(inputs) def _batch_inference(self, batched_inputs): """ Execute inference on a list of inputs, using batch size = self.batch_size, instead of the length of the list. Inputs & outputs have the same format as :meth:`SemanticSegmentor.forward` """ outputs = [] inputs = [] for idx, input in zip(count(), batched_inputs): inputs.append(input) if len(inputs) == self.batch_size or idx == len(batched_inputs) - 1: with torch.no_grad(): outputs.extend(self._inference_with_model(inputs)) inputs = [] return outputs def __call__(self, batched_inputs): """ Same input/output format as :meth:`SemanticSegmentor.forward` """ def _maybe_read_image(dataset_dict): ret = copy.copy(dataset_dict) if "image" not in ret: image = read_image(ret.pop("file_name"), self.model.input_format) image = torch.from_numpy( np.ascontiguousarray(image.transpose(2, 0, 1)) ) # CHW ret["image"] = image if "height" not in ret and "width" not in ret: ret["height"] = image.shape[1] ret["width"] = image.shape[2] return ret return [self._inference_one_image(_maybe_read_image(x)) for x in batched_inputs] def _inference_one_image(self, input): """ Args: input (dict): one dataset dict with "image" field being a CHW tensor Returns: dict: one output dict """ augmented_inputs, tfms = self._get_augmented_inputs(input) # 1: forward with all augmented images outputs = self._batch_inference(augmented_inputs) # Delete now useless variables to avoid being out of memory del augmented_inputs # 2: merge the results # handle flip specially # outputs = [output.detach() for output in outputs] return self._merge_auged_output(outputs, tfms) def _merge_auged_output(self, outputs, tfms): new_outputs = [] for output, tfm in zip(outputs, tfms): if any(isinstance(t, HFlipTransform) for t in tfm.transforms): new_outputs.append(output["sem_seg"].flip(dims=[2])) else: new_outputs.append(output["sem_seg"]) del outputs # to avoid OOM with torch.stack final_predictions = new_outputs[0] for i in range(1, len(new_outputs)): final_predictions += new_outputs[i] final_predictions = final_predictions / len(new_outputs) del new_outputs return {"sem_seg": final_predictions} def _get_augmented_inputs(self, input): augmented_inputs = self.tta_mapper(input) tfms = [x.pop("transforms") for x in augmented_inputs] return augmented_inputs, tfms