coreml-community/ControlNet-v1-1-Annotators-cpu

1# Copyright (c) Facebook, Inc. and its affiliates.2import copy3import numpy as np4from contextlib import contextmanager5from itertools import count6from typing import List7import torch8from fvcore.transforms import HFlipTransform, NoOpTransform9from torch import nn10from torch.nn.parallel import DistributedDataParallel11 12from annotator.oneformer.detectron2.config import configurable13from annotator.oneformer.detectron2.data.detection_utils import read_image14from annotator.oneformer.detectron2.data.transforms import (15    RandomFlip,16    ResizeShortestEdge,17    ResizeTransform,18    apply_augmentations,19)20from annotator.oneformer.detectron2.structures import Boxes, Instances21 22from .meta_arch import GeneralizedRCNN23from .postprocessing import detector_postprocess24from .roi_heads.fast_rcnn import fast_rcnn_inference_single_image25 26__all__ = ["DatasetMapperTTA", "GeneralizedRCNNWithTTA"]27 28 29class DatasetMapperTTA:30    """31    Implement test-time augmentation for detection data.32    It is a callable which takes a dataset dict from a detection dataset,33    and returns a list of dataset dicts where the images34    are augmented from the input image by the transformations defined in the config.35    This is used for test-time augmentation.36    """37 38    @configurable39    def __init__(self, min_sizes: List[int], max_size: int, flip: bool):40        """41        Args:42            min_sizes: list of short-edge size to resize the image to43            max_size: maximum height or width of resized images44            flip: whether to apply flipping augmentation45        """46        self.min_sizes = min_sizes47        self.max_size = max_size48        self.flip = flip49 50    @classmethod51    def from_config(cls, cfg):52        return {53            "min_sizes": cfg.TEST.AUG.MIN_SIZES,54            "max_size": cfg.TEST.AUG.MAX_SIZE,55            "flip": cfg.TEST.AUG.FLIP,56        }57 58    def __call__(self, dataset_dict):59        """60        Args:61            dict: a dict in standard model input format. See tutorials for details.62 63        Returns:64            list[dict]:65                a list of dicts, which contain augmented version of the input image.66                The total number of dicts is ``len(min_sizes) * (2 if flip else 1)``.67                Each dict has field "transforms" which is a TransformList,68                containing the transforms that are used to generate this image.69        """70        numpy_image = dataset_dict["image"].permute(1, 2, 0).numpy()71        shape = numpy_image.shape72        orig_shape = (dataset_dict["height"], dataset_dict["width"])73        if shape[:2] != orig_shape:74            # It transforms the "original" image in the dataset to the input image75            pre_tfm = ResizeTransform(orig_shape[0], orig_shape[1], shape[0], shape[1])76        else:77            pre_tfm = NoOpTransform()78 79        # Create all combinations of augmentations to use80        aug_candidates = []  # each element is a list[Augmentation]81        for min_size in self.min_sizes:82            resize = ResizeShortestEdge(min_size, self.max_size)83            aug_candidates.append([resize])  # resize only84            if self.flip:85                flip = RandomFlip(prob=1.0)86                aug_candidates.append([resize, flip])  # resize + flip87 88        # Apply all the augmentations89        ret = []90        for aug in aug_candidates:91            new_image, tfms = apply_augmentations(aug, np.copy(numpy_image))92            torch_image = torch.from_numpy(np.ascontiguousarray(new_image.transpose(2, 0, 1)))93 94            dic = copy.deepcopy(dataset_dict)95            dic["transforms"] = pre_tfm + tfms96            dic["image"] = torch_image97            ret.append(dic)98        return ret99 100 101class GeneralizedRCNNWithTTA(nn.Module):102    """103    A GeneralizedRCNN with test-time augmentation enabled.104    Its :meth:`__call__` method has the same interface as :meth:`GeneralizedRCNN.forward`.105    """106 107    def __init__(self, cfg, model, tta_mapper=None, batch_size=3):108        """109        Args:110            cfg (CfgNode):111            model (GeneralizedRCNN): a GeneralizedRCNN to apply TTA on.112            tta_mapper (callable): takes a dataset dict and returns a list of113                augmented versions of the dataset dict. Defaults to114                `DatasetMapperTTA(cfg)`.115            batch_size (int): batch the augmented images into this batch size for inference.116        """117        super().__init__()118        if isinstance(model, DistributedDataParallel):119            model = model.module120        assert isinstance(121            model, GeneralizedRCNN122        ), "TTA is only supported on GeneralizedRCNN. Got a model of type {}".format(type(model))123        self.cfg = cfg.clone()124        assert not self.cfg.MODEL.KEYPOINT_ON, "TTA for keypoint is not supported yet"125        assert (126            not self.cfg.MODEL.LOAD_PROPOSALS127        ), "TTA for pre-computed proposals is not supported yet"128 129        self.model = model130 131        if tta_mapper is None:132            tta_mapper = DatasetMapperTTA(cfg)133        self.tta_mapper = tta_mapper134        self.batch_size = batch_size135 136    @contextmanager137    def _turn_off_roi_heads(self, attrs):138        """139        Open a context where some heads in `model.roi_heads` are temporarily turned off.140        Args:141            attr (list[str]): the attribute in `model.roi_heads` which can be used142                to turn off a specific head, e.g., "mask_on", "keypoint_on".143        """144        roi_heads = self.model.roi_heads145        old = {}146        for attr in attrs:147            try:148                old[attr] = getattr(roi_heads, attr)149            except AttributeError:150                # The head may not be implemented in certain ROIHeads151                pass152 153        if len(old.keys()) == 0:154            yield155        else:156            for attr in old.keys():157                setattr(roi_heads, attr, False)158            yield159            for attr in old.keys():160                setattr(roi_heads, attr, old[attr])161 162    def _batch_inference(self, batched_inputs, detected_instances=None):163        """164        Execute inference on a list of inputs,165        using batch size = self.batch_size, instead of the length of the list.166 167        Inputs & outputs have the same format as :meth:`GeneralizedRCNN.inference`168        """169        if detected_instances is None:170            detected_instances = [None] * len(batched_inputs)171 172        outputs = []173        inputs, instances = [], []174        for idx, input, instance in zip(count(), batched_inputs, detected_instances):175            inputs.append(input)176            instances.append(instance)177            if len(inputs) == self.batch_size or idx == len(batched_inputs) - 1:178                outputs.extend(179                    self.model.inference(180                        inputs,181                        instances if instances[0] is not None else None,182                        do_postprocess=False,183                    )184                )185                inputs, instances = [], []186        return outputs187 188    def __call__(self, batched_inputs):189        """190        Same input/output format as :meth:`GeneralizedRCNN.forward`191        """192 193        def _maybe_read_image(dataset_dict):194            ret = copy.copy(dataset_dict)195            if "image" not in ret:196                image = read_image(ret.pop("file_name"), self.model.input_format)197                image = torch.from_numpy(np.ascontiguousarray(image.transpose(2, 0, 1)))  # CHW198                ret["image"] = image199            if "height" not in ret and "width" not in ret:200                ret["height"] = image.shape[1]201                ret["width"] = image.shape[2]202            return ret203 204        return [self._inference_one_image(_maybe_read_image(x)) for x in batched_inputs]205 206    def _inference_one_image(self, input):207        """208        Args:209            input (dict): one dataset dict with "image" field being a CHW tensor210 211        Returns:212            dict: one output dict213        """214        orig_shape = (input["height"], input["width"])215        augmented_inputs, tfms = self._get_augmented_inputs(input)216        # Detect boxes from all augmented versions217        with self._turn_off_roi_heads(["mask_on", "keypoint_on"]):218            # temporarily disable roi heads219            all_boxes, all_scores, all_classes = self._get_augmented_boxes(augmented_inputs, tfms)220        # merge all detected boxes to obtain final predictions for boxes221        merged_instances = self._merge_detections(all_boxes, all_scores, all_classes, orig_shape)222 223        if self.cfg.MODEL.MASK_ON:224            # Use the detected boxes to obtain masks225            augmented_instances = self._rescale_detected_boxes(226                augmented_inputs, merged_instances, tfms227            )228            # run forward on the detected boxes229            outputs = self._batch_inference(augmented_inputs, augmented_instances)230            # Delete now useless variables to avoid being out of memory231            del augmented_inputs, augmented_instances232            # average the predictions233            merged_instances.pred_masks = self._reduce_pred_masks(outputs, tfms)234            merged_instances = detector_postprocess(merged_instances, *orig_shape)235            return {"instances": merged_instances}236        else:237            return {"instances": merged_instances}238 239    def _get_augmented_inputs(self, input):240        augmented_inputs = self.tta_mapper(input)241        tfms = [x.pop("transforms") for x in augmented_inputs]242        return augmented_inputs, tfms243 244    def _get_augmented_boxes(self, augmented_inputs, tfms):245        # 1: forward with all augmented images246        outputs = self._batch_inference(augmented_inputs)247        # 2: union the results248        all_boxes = []249        all_scores = []250        all_classes = []251        for output, tfm in zip(outputs, tfms):252            # Need to inverse the transforms on boxes, to obtain results on original image253            pred_boxes = output.pred_boxes.tensor254            original_pred_boxes = tfm.inverse().apply_box(pred_boxes.cpu().numpy())255            all_boxes.append(torch.from_numpy(original_pred_boxes).to(pred_boxes.device))256 257            all_scores.extend(output.scores)258            all_classes.extend(output.pred_classes)259        all_boxes = torch.cat(all_boxes, dim=0)260        return all_boxes, all_scores, all_classes261 262    def _merge_detections(self, all_boxes, all_scores, all_classes, shape_hw):263        # select from the union of all results264        num_boxes = len(all_boxes)265        num_classes = self.cfg.MODEL.ROI_HEADS.NUM_CLASSES266        # +1 because fast_rcnn_inference expects background scores as well267        all_scores_2d = torch.zeros(num_boxes, num_classes + 1, device=all_boxes.device)268        for idx, cls, score in zip(count(), all_classes, all_scores):269            all_scores_2d[idx, cls] = score270 271        merged_instances, _ = fast_rcnn_inference_single_image(272            all_boxes,273            all_scores_2d,274            shape_hw,275            1e-8,276            self.cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST,277            self.cfg.TEST.DETECTIONS_PER_IMAGE,278        )279 280        return merged_instances281 282    def _rescale_detected_boxes(self, augmented_inputs, merged_instances, tfms):283        augmented_instances = []284        for input, tfm in zip(augmented_inputs, tfms):285            # Transform the target box to the augmented image's coordinate space286            pred_boxes = merged_instances.pred_boxes.tensor.cpu().numpy()287            pred_boxes = torch.from_numpy(tfm.apply_box(pred_boxes))288 289            aug_instances = Instances(290                image_size=input["image"].shape[1:3],291                pred_boxes=Boxes(pred_boxes),292                pred_classes=merged_instances.pred_classes,293                scores=merged_instances.scores,294            )295            augmented_instances.append(aug_instances)296        return augmented_instances297 298    def _reduce_pred_masks(self, outputs, tfms):299        # Should apply inverse transforms on masks.300        # We assume only resize & flip are used. pred_masks is a scale-invariant301        # representation, so we handle flip specially302        for output, tfm in zip(outputs, tfms):303            if any(isinstance(t, HFlipTransform) for t in tfm.transforms):304                output.pred_masks = output.pred_masks.flip(dims=[3])305        all_pred_masks = torch.stack([o.pred_masks for o in outputs], dim=0)306        avg_pred_masks = torch.mean(all_pred_masks, dim=0)307        return avg_pred_masks308