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

1# Copyright (c) Facebook, Inc. and its affiliates.2import itertools3import logging4import numpy as np5from collections import OrderedDict6from collections.abc import Mapping7from typing import Dict, List, Optional, Tuple, Union8import torch9from omegaconf import DictConfig, OmegaConf10from torch import Tensor, nn11 12from annotator.oneformer.detectron2.layers import ShapeSpec13from annotator.oneformer.detectron2.structures import BitMasks, Boxes, ImageList, Instances14from annotator.oneformer.detectron2.utils.events import get_event_storage15 16from .backbone import Backbone17 18logger = logging.getLogger(__name__)19 20 21def _to_container(cfg):22    """23    mmdet will assert the type of dict/list.24    So convert omegaconf objects to dict/list.25    """26    if isinstance(cfg, DictConfig):27        cfg = OmegaConf.to_container(cfg, resolve=True)28    from mmcv.utils import ConfigDict29 30    return ConfigDict(cfg)31 32 33class MMDetBackbone(Backbone):34    """35    Wrapper of mmdetection backbones to use in detectron2.36 37    mmdet backbones produce list/tuple of tensors, while detectron2 backbones38    produce a dict of tensors. This class wraps the given backbone to produce39    output in detectron2's convention, so it can be used in place of detectron240    backbones.41    """42 43    def __init__(44        self,45        backbone: Union[nn.Module, Mapping],46        neck: Union[nn.Module, Mapping, None] = None,47        *,48        output_shapes: List[ShapeSpec],49        output_names: Optional[List[str]] = None,50    ):51        """52        Args:53            backbone: either a backbone module or a mmdet config dict that defines a54                backbone. The backbone takes a 4D image tensor and returns a55                sequence of tensors.56            neck: either a backbone module or a mmdet config dict that defines a57                neck. The neck takes outputs of backbone and returns a58                sequence of tensors. If None, no neck is used.59            output_shapes: shape for every output of the backbone (or neck, if given).60                stride and channels are often needed.61            output_names: names for every output of the backbone (or neck, if given).62                By default, will use "out0", "out1", ...63        """64        super().__init__()65        if isinstance(backbone, Mapping):66            from mmdet.models import build_backbone67 68            backbone = build_backbone(_to_container(backbone))69        self.backbone = backbone70 71        if isinstance(neck, Mapping):72            from mmdet.models import build_neck73 74            neck = build_neck(_to_container(neck))75        self.neck = neck76 77        # "Neck" weights, if any, are part of neck itself. This is the interface78        # of mmdet so we follow it. Reference:79        # https://github.com/open-mmlab/mmdetection/blob/master/mmdet/models/detectors/two_stage.py80        logger.info("Initializing mmdet backbone weights...")81        self.backbone.init_weights()82        # train() in mmdet modules is non-trivial, and has to be explicitly83        # called. Reference:84        # https://github.com/open-mmlab/mmdetection/blob/master/mmdet/models/backbones/resnet.py85        self.backbone.train()86        if self.neck is not None:87            logger.info("Initializing mmdet neck weights ...")88            if isinstance(self.neck, nn.Sequential):89                for m in self.neck:90                    m.init_weights()91            else:92                self.neck.init_weights()93            self.neck.train()94 95        self._output_shapes = output_shapes96        if not output_names:97            output_names = [f"out{i}" for i in range(len(output_shapes))]98        self._output_names = output_names99 100    def forward(self, x) -> Dict[str, Tensor]:101        outs = self.backbone(x)102        if self.neck is not None:103            outs = self.neck(outs)104        assert isinstance(105            outs, (list, tuple)106        ), "mmdet backbone should return a list/tuple of tensors!"107        if len(outs) != len(self._output_shapes):108            raise ValueError(109                "Length of output_shapes does not match outputs from the mmdet backbone: "110                f"{len(outs)} != {len(self._output_shapes)}"111            )112        return {k: v for k, v in zip(self._output_names, outs)}113 114    def output_shape(self) -> Dict[str, ShapeSpec]:115        return {k: v for k, v in zip(self._output_names, self._output_shapes)}116 117 118class MMDetDetector(nn.Module):119    """120    Wrapper of a mmdetection detector model, for detection and instance segmentation.121    Input/output formats of this class follow detectron2's convention, so a122    mmdetection model can be trained and evaluated in detectron2.123    """124 125    def __init__(126        self,127        detector: Union[nn.Module, Mapping],128        *,129        # Default is 32 regardless of model:130        # https://github.com/open-mmlab/mmdetection/tree/master/configs/_base_/datasets131        size_divisibility=32,132        pixel_mean: Tuple[float],133        pixel_std: Tuple[float],134    ):135        """136        Args:137            detector: a mmdet detector, or a mmdet config dict that defines a detector.138            size_divisibility: pad input images to multiple of this number139            pixel_mean: per-channel mean to normalize input image140            pixel_std: per-channel stddev to normalize input image141        """142        super().__init__()143        if isinstance(detector, Mapping):144            from mmdet.models import build_detector145 146            detector = build_detector(_to_container(detector))147        self.detector = detector148        self.detector.init_weights()149        self.size_divisibility = size_divisibility150 151        self.register_buffer("pixel_mean", torch.tensor(pixel_mean).view(-1, 1, 1), False)152        self.register_buffer("pixel_std", torch.tensor(pixel_std).view(-1, 1, 1), False)153        assert (154            self.pixel_mean.shape == self.pixel_std.shape155        ), f"{self.pixel_mean} and {self.pixel_std} have different shapes!"156 157    def forward(self, batched_inputs: List[Dict[str, torch.Tensor]]):158        images = [x["image"].to(self.device) for x in batched_inputs]159        images = [(x - self.pixel_mean) / self.pixel_std for x in images]160        images = ImageList.from_tensors(images, size_divisibility=self.size_divisibility).tensor161        metas = []162        rescale = {"height" in x for x in batched_inputs}163        if len(rescale) != 1:164            raise ValueError("Some inputs have original height/width, but some don't!")165        rescale = list(rescale)[0]166        output_shapes = []167        for input in batched_inputs:168            meta = {}169            c, h, w = input["image"].shape170            meta["img_shape"] = meta["ori_shape"] = (h, w, c)171            if rescale:172                scale_factor = np.array(173                    [w / input["width"], h / input["height"]] * 2, dtype="float32"174                )175                ori_shape = (input["height"], input["width"])176                output_shapes.append(ori_shape)177                meta["ori_shape"] = ori_shape + (c,)178            else:179                scale_factor = 1.0180                output_shapes.append((h, w))181            meta["scale_factor"] = scale_factor182            meta["flip"] = False183            padh, padw = images.shape[-2:]184            meta["pad_shape"] = (padh, padw, c)185            metas.append(meta)186 187        if self.training:188            gt_instances = [x["instances"].to(self.device) for x in batched_inputs]189            if gt_instances[0].has("gt_masks"):190                from mmdet.core import PolygonMasks as mm_PolygonMasks, BitmapMasks as mm_BitMasks191 192                def convert_mask(m, shape):193                    # mmdet mask format194                    if isinstance(m, BitMasks):195                        return mm_BitMasks(m.tensor.cpu().numpy(), shape[0], shape[1])196                    else:197                        return mm_PolygonMasks(m.polygons, shape[0], shape[1])198 199                gt_masks = [convert_mask(x.gt_masks, x.image_size) for x in gt_instances]200                losses_and_metrics = self.detector.forward_train(201                    images,202                    metas,203                    [x.gt_boxes.tensor for x in gt_instances],204                    [x.gt_classes for x in gt_instances],205                    gt_masks=gt_masks,206                )207            else:208                losses_and_metrics = self.detector.forward_train(209                    images,210                    metas,211                    [x.gt_boxes.tensor for x in gt_instances],212                    [x.gt_classes for x in gt_instances],213                )214            return _parse_losses(losses_and_metrics)215        else:216            results = self.detector.simple_test(images, metas, rescale=rescale)217            results = [218                {"instances": _convert_mmdet_result(r, shape)}219                for r, shape in zip(results, output_shapes)220            ]221            return results222 223    @property224    def device(self):225        return self.pixel_mean.device226 227 228# Reference: show_result() in229# https://github.com/open-mmlab/mmdetection/blob/master/mmdet/models/detectors/base.py230def _convert_mmdet_result(result, shape: Tuple[int, int]) -> Instances:231    if isinstance(result, tuple):232        bbox_result, segm_result = result233        if isinstance(segm_result, tuple):234            segm_result = segm_result[0]235    else:236        bbox_result, segm_result = result, None237 238    bboxes = torch.from_numpy(np.vstack(bbox_result))  # Nx5239    bboxes, scores = bboxes[:, :4], bboxes[:, -1]240    labels = [241        torch.full((bbox.shape[0],), i, dtype=torch.int32) for i, bbox in enumerate(bbox_result)242    ]243    labels = torch.cat(labels)244    inst = Instances(shape)245    inst.pred_boxes = Boxes(bboxes)246    inst.scores = scores247    inst.pred_classes = labels248 249    if segm_result is not None and len(labels) > 0:250        segm_result = list(itertools.chain(*segm_result))251        segm_result = [torch.from_numpy(x) if isinstance(x, np.ndarray) else x for x in segm_result]252        segm_result = torch.stack(segm_result, dim=0)253        inst.pred_masks = segm_result254    return inst255 256 257# reference: https://github.com/open-mmlab/mmdetection/blob/master/mmdet/models/detectors/base.py258def _parse_losses(losses: Dict[str, Tensor]) -> Dict[str, Tensor]:259    log_vars = OrderedDict()260    for loss_name, loss_value in losses.items():261        if isinstance(loss_value, torch.Tensor):262            log_vars[loss_name] = loss_value.mean()263        elif isinstance(loss_value, list):264            log_vars[loss_name] = sum(_loss.mean() for _loss in loss_value)265        else:266            raise TypeError(f"{loss_name} is not a tensor or list of tensors")267 268        if "loss" not in loss_name:269            # put metrics to storage; don't return them270            storage = get_event_storage()271            value = log_vars.pop(loss_name).cpu().item()272            storage.put_scalar(loss_name, value)273    return log_vars274