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

1# Copyright (c) Facebook, Inc. and its affiliates.2import copy3import logging4import numpy as np5from typing import List, Optional, Union6import torch7 8from annotator.oneformer.detectron2.config import configurable9 10from . import detection_utils as utils11from . import transforms as T12 13"""14This file contains the default mapping that's applied to "dataset dicts".15"""16 17__all__ = ["DatasetMapper"]18 19 20class DatasetMapper:21    """22    A callable which takes a dataset dict in Detectron2 Dataset format,23    and map it into a format used by the model.24 25    This is the default callable to be used to map your dataset dict into training data.26    You may need to follow it to implement your own one for customized logic,27    such as a different way to read or transform images.28    See :doc:`/tutorials/data_loading` for details.29 30    The callable currently does the following:31 32    1. Read the image from "file_name"33    2. Applies cropping/geometric transforms to the image and annotations34    3. Prepare data and annotations to Tensor and :class:`Instances`35    """36 37    @configurable38    def __init__(39        self,40        is_train: bool,41        *,42        augmentations: List[Union[T.Augmentation, T.Transform]],43        image_format: str,44        use_instance_mask: bool = False,45        use_keypoint: bool = False,46        instance_mask_format: str = "polygon",47        keypoint_hflip_indices: Optional[np.ndarray] = None,48        precomputed_proposal_topk: Optional[int] = None,49        recompute_boxes: bool = False,50    ):51        """52        NOTE: this interface is experimental.53 54        Args:55            is_train: whether it's used in training or inference56            augmentations: a list of augmentations or deterministic transforms to apply57            image_format: an image format supported by :func:`detection_utils.read_image`.58            use_instance_mask: whether to process instance segmentation annotations, if available59            use_keypoint: whether to process keypoint annotations if available60            instance_mask_format: one of "polygon" or "bitmask". Process instance segmentation61                masks into this format.62            keypoint_hflip_indices: see :func:`detection_utils.create_keypoint_hflip_indices`63            precomputed_proposal_topk: if given, will load pre-computed64                proposals from dataset_dict and keep the top k proposals for each image.65            recompute_boxes: whether to overwrite bounding box annotations66                by computing tight bounding boxes from instance mask annotations.67        """68        if recompute_boxes:69            assert use_instance_mask, "recompute_boxes requires instance masks"70        # fmt: off71        self.is_train               = is_train72        self.augmentations          = T.AugmentationList(augmentations)73        self.image_format           = image_format74        self.use_instance_mask      = use_instance_mask75        self.instance_mask_format   = instance_mask_format76        self.use_keypoint           = use_keypoint77        self.keypoint_hflip_indices = keypoint_hflip_indices78        self.proposal_topk          = precomputed_proposal_topk79        self.recompute_boxes        = recompute_boxes80        # fmt: on81        logger = logging.getLogger(__name__)82        mode = "training" if is_train else "inference"83        logger.info(f"[DatasetMapper] Augmentations used in {mode}: {augmentations}")84 85    @classmethod86    def from_config(cls, cfg, is_train: bool = True):87        augs = utils.build_augmentation(cfg, is_train)88        if cfg.INPUT.CROP.ENABLED and is_train:89            augs.insert(0, T.RandomCrop(cfg.INPUT.CROP.TYPE, cfg.INPUT.CROP.SIZE))90            recompute_boxes = cfg.MODEL.MASK_ON91        else:92            recompute_boxes = False93 94        ret = {95            "is_train": is_train,96            "augmentations": augs,97            "image_format": cfg.INPUT.FORMAT,98            "use_instance_mask": cfg.MODEL.MASK_ON,99            "instance_mask_format": cfg.INPUT.MASK_FORMAT,100            "use_keypoint": cfg.MODEL.KEYPOINT_ON,101            "recompute_boxes": recompute_boxes,102        }103 104        if cfg.MODEL.KEYPOINT_ON:105            ret["keypoint_hflip_indices"] = utils.create_keypoint_hflip_indices(cfg.DATASETS.TRAIN)106 107        if cfg.MODEL.LOAD_PROPOSALS:108            ret["precomputed_proposal_topk"] = (109                cfg.DATASETS.PRECOMPUTED_PROPOSAL_TOPK_TRAIN110                if is_train111                else cfg.DATASETS.PRECOMPUTED_PROPOSAL_TOPK_TEST112            )113        return ret114 115    def _transform_annotations(self, dataset_dict, transforms, image_shape):116        # USER: Modify this if you want to keep them for some reason.117        for anno in dataset_dict["annotations"]:118            if not self.use_instance_mask:119                anno.pop("segmentation", None)120            if not self.use_keypoint:121                anno.pop("keypoints", None)122 123        # USER: Implement additional transformations if you have other types of data124        annos = [125            utils.transform_instance_annotations(126                obj, transforms, image_shape, keypoint_hflip_indices=self.keypoint_hflip_indices127            )128            for obj in dataset_dict.pop("annotations")129            if obj.get("iscrowd", 0) == 0130        ]131        instances = utils.annotations_to_instances(132            annos, image_shape, mask_format=self.instance_mask_format133        )134 135        # After transforms such as cropping are applied, the bounding box may no longer136        # tightly bound the object. As an example, imagine a triangle object137        # [(0,0), (2,0), (0,2)] cropped by a box [(1,0),(2,2)] (XYXY format). The tight138        # bounding box of the cropped triangle should be [(1,0),(2,1)], which is not equal to139        # the intersection of original bounding box and the cropping box.140        if self.recompute_boxes:141            instances.gt_boxes = instances.gt_masks.get_bounding_boxes()142        dataset_dict["instances"] = utils.filter_empty_instances(instances)143 144    def __call__(self, dataset_dict):145        """146        Args:147            dataset_dict (dict): Metadata of one image, in Detectron2 Dataset format.148 149        Returns:150            dict: a format that builtin models in detectron2 accept151        """152        dataset_dict = copy.deepcopy(dataset_dict)  # it will be modified by code below153        # USER: Write your own image loading if it's not from a file154        image = utils.read_image(dataset_dict["file_name"], format=self.image_format)155        utils.check_image_size(dataset_dict, image)156 157        # USER: Remove if you don't do semantic/panoptic segmentation.158        if "sem_seg_file_name" in dataset_dict:159            sem_seg_gt = utils.read_image(dataset_dict.pop("sem_seg_file_name"), "L").squeeze(2)160        else:161            sem_seg_gt = None162 163        aug_input = T.AugInput(image, sem_seg=sem_seg_gt)164        transforms = self.augmentations(aug_input)165        image, sem_seg_gt = aug_input.image, aug_input.sem_seg166 167        image_shape = image.shape[:2]  # h, w168        # Pytorch's dataloader is efficient on torch.Tensor due to shared-memory,169        # but not efficient on large generic data structures due to the use of pickle & mp.Queue.170        # Therefore it's important to use torch.Tensor.171        dataset_dict["image"] = torch.as_tensor(np.ascontiguousarray(image.transpose(2, 0, 1)))172        if sem_seg_gt is not None:173            dataset_dict["sem_seg"] = torch.as_tensor(sem_seg_gt.astype("long"))174 175        # USER: Remove if you don't use pre-computed proposals.176        # Most users would not need this feature.177        if self.proposal_topk is not None:178            utils.transform_proposals(179                dataset_dict, image_shape, transforms, proposal_topk=self.proposal_topk180            )181 182        if not self.is_train:183            # USER: Modify this if you want to keep them for some reason.184            dataset_dict.pop("annotations", None)185            dataset_dict.pop("sem_seg_file_name", None)186            return dataset_dict187 188        if "annotations" in dataset_dict:189            self._transform_annotations(dataset_dict, transforms, image_shape)190 191        return dataset_dict192