coreml-community/ControlNet-v1-1-Annotators-cpu
15
1# -*- coding: utf-8 -*-2# Copyright (c) Facebook, Inc. and its affiliates.3 4"""5Common data processing utilities that are used in a6typical object detection data pipeline.7"""8import logging9import numpy as np10from typing import List, Union11import pycocotools.mask as mask_util12import torch13from PIL import Image14 15from annotator.oneformer.detectron2.structures import (16 BitMasks,17 Boxes,18 BoxMode,19 Instances,20 Keypoints,21 PolygonMasks,22 RotatedBoxes,23 polygons_to_bitmask,24)25from annotator.oneformer.detectron2.utils.file_io import PathManager26 27from . import transforms as T28from .catalog import MetadataCatalog29 30__all__ = [31 "SizeMismatchError",32 "convert_image_to_rgb",33 "check_image_size",34 "transform_proposals",35 "transform_instance_annotations",36 "annotations_to_instances",37 "annotations_to_instances_rotated",38 "build_augmentation",39 "build_transform_gen",40 "create_keypoint_hflip_indices",41 "filter_empty_instances",42 "read_image",43]44 45 46class SizeMismatchError(ValueError):47 """48 When loaded image has difference width/height compared with annotation.49 """50 51 52# https://en.wikipedia.org/wiki/YUV#SDTV_with_BT.60153_M_RGB2YUV = [[0.299, 0.587, 0.114], [-0.14713, -0.28886, 0.436], [0.615, -0.51499, -0.10001]]54_M_YUV2RGB = [[1.0, 0.0, 1.13983], [1.0, -0.39465, -0.58060], [1.0, 2.03211, 0.0]]55 56# https://www.exiv2.org/tags.html57_EXIF_ORIENT = 274 # exif 'Orientation' tag58 59 60def convert_PIL_to_numpy(image, format):61 """62 Convert PIL image to numpy array of target format.63 64 Args:65 image (PIL.Image): a PIL image66 format (str): the format of output image67 68 Returns:69 (np.ndarray): also see `read_image`70 """71 if format is not None:72 # PIL only supports RGB, so convert to RGB and flip channels over below73 conversion_format = format74 if format in ["BGR", "YUV-BT.601"]:75 conversion_format = "RGB"76 image = image.convert(conversion_format)77 image = np.asarray(image)78 # PIL squeezes out the channel dimension for "L", so make it HWC79 if format == "L":80 image = np.expand_dims(image, -1)81 82 # handle formats not supported by PIL83 elif format == "BGR":84 # flip channels if needed85 image = image[:, :, ::-1]86 elif format == "YUV-BT.601":87 image = image / 255.088 image = np.dot(image, np.array(_M_RGB2YUV).T)89 90 return image91 92 93def convert_image_to_rgb(image, format):94 """95 Convert an image from given format to RGB.96 97 Args:98 image (np.ndarray or Tensor): an HWC image99 format (str): the format of input image, also see `read_image`100 101 Returns:102 (np.ndarray): (H,W,3) RGB image in 0-255 range, can be either float or uint8103 """104 if isinstance(image, torch.Tensor):105 image = image.cpu().numpy()106 if format == "BGR":107 image = image[:, :, [2, 1, 0]]108 elif format == "YUV-BT.601":109 image = np.dot(image, np.array(_M_YUV2RGB).T)110 image = image * 255.0111 else:112 if format == "L":113 image = image[:, :, 0]114 image = image.astype(np.uint8)115 image = np.asarray(Image.fromarray(image, mode=format).convert("RGB"))116 return image117 118 119def _apply_exif_orientation(image):120 """121 Applies the exif orientation correctly.122 123 This code exists per the bug:124 https://github.com/python-pillow/Pillow/issues/3973125 with the function `ImageOps.exif_transpose`. The Pillow source raises errors with126 various methods, especially `tobytes`127 128 Function based on:129 https://github.com/wkentaro/labelme/blob/v4.5.4/labelme/utils/image.py#L59130 https://github.com/python-pillow/Pillow/blob/7.1.2/src/PIL/ImageOps.py#L527131 132 Args:133 image (PIL.Image): a PIL image134 135 Returns:136 (PIL.Image): the PIL image with exif orientation applied, if applicable137 """138 if not hasattr(image, "getexif"):139 return image140 141 try:142 exif = image.getexif()143 except Exception: # https://github.com/facebookresearch/detectron2/issues/1885144 exif = None145 146 if exif is None:147 return image148 149 orientation = exif.get(_EXIF_ORIENT)150 151 method = {152 2: Image.FLIP_LEFT_RIGHT,153 3: Image.ROTATE_180,154 4: Image.FLIP_TOP_BOTTOM,155 5: Image.TRANSPOSE,156 6: Image.ROTATE_270,157 7: Image.TRANSVERSE,158 8: Image.ROTATE_90,159 }.get(orientation)160 161 if method is not None:162 return image.transpose(method)163 return image164 165 166def read_image(file_name, format=None):167 """168 Read an image into the given format.169 Will apply rotation and flipping if the image has such exif information.170 171 Args:172 file_name (str): image file path173 format (str): one of the supported image modes in PIL, or "BGR" or "YUV-BT.601".174 175 Returns:176 image (np.ndarray):177 an HWC image in the given format, which is 0-255, uint8 for178 supported image modes in PIL or "BGR"; float (0-1 for Y) for YUV-BT.601.179 """180 with PathManager.open(file_name, "rb") as f:181 image = Image.open(f)182 183 # work around this bug: https://github.com/python-pillow/Pillow/issues/3973184 image = _apply_exif_orientation(image)185 return convert_PIL_to_numpy(image, format)186 187 188def check_image_size(dataset_dict, image):189 """190 Raise an error if the image does not match the size specified in the dict.191 """192 if "width" in dataset_dict or "height" in dataset_dict:193 image_wh = (image.shape[1], image.shape[0])194 expected_wh = (dataset_dict["width"], dataset_dict["height"])195 if not image_wh == expected_wh:196 raise SizeMismatchError(197 "Mismatched image shape{}, got {}, expect {}.".format(198 " for image " + dataset_dict["file_name"]199 if "file_name" in dataset_dict200 else "",201 image_wh,202 expected_wh,203 )204 + " Please check the width/height in your annotation."205 )206 207 # To ensure bbox always remap to original image size208 if "width" not in dataset_dict:209 dataset_dict["width"] = image.shape[1]210 if "height" not in dataset_dict:211 dataset_dict["height"] = image.shape[0]212 213 214def transform_proposals(dataset_dict, image_shape, transforms, *, proposal_topk, min_box_size=0):215 """216 Apply transformations to the proposals in dataset_dict, if any.217 218 Args:219 dataset_dict (dict): a dict read from the dataset, possibly220 contains fields "proposal_boxes", "proposal_objectness_logits", "proposal_bbox_mode"221 image_shape (tuple): height, width222 transforms (TransformList):223 proposal_topk (int): only keep top-K scoring proposals224 min_box_size (int): proposals with either side smaller than this225 threshold are removed226 227 The input dict is modified in-place, with abovementioned keys removed. A new228 key "proposals" will be added. Its value is an `Instances`229 object which contains the transformed proposals in its field230 "proposal_boxes" and "objectness_logits".231 """232 if "proposal_boxes" in dataset_dict:233 # Transform proposal boxes234 boxes = transforms.apply_box(235 BoxMode.convert(236 dataset_dict.pop("proposal_boxes"),237 dataset_dict.pop("proposal_bbox_mode"),238 BoxMode.XYXY_ABS,239 )240 )241 boxes = Boxes(boxes)242 objectness_logits = torch.as_tensor(243 dataset_dict.pop("proposal_objectness_logits").astype("float32")244 )245 246 boxes.clip(image_shape)247 keep = boxes.nonempty(threshold=min_box_size)248 boxes = boxes[keep]249 objectness_logits = objectness_logits[keep]250 251 proposals = Instances(image_shape)252 proposals.proposal_boxes = boxes[:proposal_topk]253 proposals.objectness_logits = objectness_logits[:proposal_topk]254 dataset_dict["proposals"] = proposals255 256 257def get_bbox(annotation):258 """259 Get bbox from data260 Args:261 annotation (dict): dict of instance annotations for a single instance.262 Returns:263 bbox (ndarray): x1, y1, x2, y2 coordinates264 """265 # bbox is 1d (per-instance bounding box)266 bbox = BoxMode.convert(annotation["bbox"], annotation["bbox_mode"], BoxMode.XYXY_ABS)267 return bbox268 269 270def transform_instance_annotations(271 annotation, transforms, image_size, *, keypoint_hflip_indices=None272):273 """274 Apply transforms to box, segmentation and keypoints annotations of a single instance.275 276 It will use `transforms.apply_box` for the box, and277 `transforms.apply_coords` for segmentation polygons & keypoints.278 If you need anything more specially designed for each data structure,279 you'll need to implement your own version of this function or the transforms.280 281 Args:282 annotation (dict): dict of instance annotations for a single instance.283 It will be modified in-place.284 transforms (TransformList or list[Transform]):285 image_size (tuple): the height, width of the transformed image286 keypoint_hflip_indices (ndarray[int]): see `create_keypoint_hflip_indices`.287 288 Returns:289 dict:290 the same input dict with fields "bbox", "segmentation", "keypoints"291 transformed according to `transforms`.292 The "bbox_mode" field will be set to XYXY_ABS.293 """294 if isinstance(transforms, (tuple, list)):295 transforms = T.TransformList(transforms)296 # bbox is 1d (per-instance bounding box)297 bbox = BoxMode.convert(annotation["bbox"], annotation["bbox_mode"], BoxMode.XYXY_ABS)298 # clip transformed bbox to image size299 bbox = transforms.apply_box(np.array([bbox]))[0].clip(min=0)300 annotation["bbox"] = np.minimum(bbox, list(image_size + image_size)[::-1])301 annotation["bbox_mode"] = BoxMode.XYXY_ABS302 303 if "segmentation" in annotation:304 # each instance contains 1 or more polygons305 segm = annotation["segmentation"]306 if isinstance(segm, list):307 # polygons308 polygons = [np.asarray(p).reshape(-1, 2) for p in segm]309 annotation["segmentation"] = [310 p.reshape(-1) for p in transforms.apply_polygons(polygons)311 ]312 elif isinstance(segm, dict):313 # RLE314 mask = mask_util.decode(segm)315 mask = transforms.apply_segmentation(mask)316 assert tuple(mask.shape[:2]) == image_size317 annotation["segmentation"] = mask318 else:319 raise ValueError(320 "Cannot transform segmentation of type '{}'!"321 "Supported types are: polygons as list[list[float] or ndarray],"322 " COCO-style RLE as a dict.".format(type(segm))323 )324 325 if "keypoints" in annotation:326 keypoints = transform_keypoint_annotations(327 annotation["keypoints"], transforms, image_size, keypoint_hflip_indices328 )329 annotation["keypoints"] = keypoints330 331 return annotation332 333 334def transform_keypoint_annotations(keypoints, transforms, image_size, keypoint_hflip_indices=None):335 """336 Transform keypoint annotations of an image.337 If a keypoint is transformed out of image boundary, it will be marked "unlabeled" (visibility=0)338 339 Args:340 keypoints (list[float]): Nx3 float in Detectron2's Dataset format.341 Each point is represented by (x, y, visibility).342 transforms (TransformList):343 image_size (tuple): the height, width of the transformed image344 keypoint_hflip_indices (ndarray[int]): see `create_keypoint_hflip_indices`.345 When `transforms` includes horizontal flip, will use the index346 mapping to flip keypoints.347 """348 # (N*3,) -> (N, 3)349 keypoints = np.asarray(keypoints, dtype="float64").reshape(-1, 3)350 keypoints_xy = transforms.apply_coords(keypoints[:, :2])351 352 # Set all out-of-boundary points to "unlabeled"353 inside = (keypoints_xy >= np.array([0, 0])) & (keypoints_xy <= np.array(image_size[::-1]))354 inside = inside.all(axis=1)355 keypoints[:, :2] = keypoints_xy356 keypoints[:, 2][~inside] = 0357 358 # This assumes that HorizFlipTransform is the only one that does flip359 do_hflip = sum(isinstance(t, T.HFlipTransform) for t in transforms.transforms) % 2 == 1360 361 # Alternative way: check if probe points was horizontally flipped.362 # probe = np.asarray([[0.0, 0.0], [image_width, 0.0]])363 # probe_aug = transforms.apply_coords(probe.copy())364 # do_hflip = np.sign(probe[1][0] - probe[0][0]) != np.sign(probe_aug[1][0] - probe_aug[0][0]) # noqa365 366 # If flipped, swap each keypoint with its opposite-handed equivalent367 if do_hflip:368 if keypoint_hflip_indices is None:369 raise ValueError("Cannot flip keypoints without providing flip indices!")370 if len(keypoints) != len(keypoint_hflip_indices):371 raise ValueError(372 "Keypoint data has {} points, but metadata "373 "contains {} points!".format(len(keypoints), len(keypoint_hflip_indices))374 )375 keypoints = keypoints[np.asarray(keypoint_hflip_indices, dtype=np.int32), :]376 377 # Maintain COCO convention that if visibility == 0 (unlabeled), then x, y = 0378 keypoints[keypoints[:, 2] == 0] = 0379 return keypoints380 381 382def annotations_to_instances(annos, image_size, mask_format="polygon"):383 """384 Create an :class:`Instances` object used by the models,385 from instance annotations in the dataset dict.386 387 Args:388 annos (list[dict]): a list of instance annotations in one image, each389 element for one instance.390 image_size (tuple): height, width391 392 Returns:393 Instances:394 It will contain fields "gt_boxes", "gt_classes",395 "gt_masks", "gt_keypoints", if they can be obtained from `annos`.396 This is the format that builtin models expect.397 """398 boxes = (399 np.stack(400 [BoxMode.convert(obj["bbox"], obj["bbox_mode"], BoxMode.XYXY_ABS) for obj in annos]401 )402 if len(annos)403 else np.zeros((0, 4))404 )405 target = Instances(image_size)406 target.gt_boxes = Boxes(boxes)407 408 classes = [int(obj["category_id"]) for obj in annos]409 classes = torch.tensor(classes, dtype=torch.int64)410 target.gt_classes = classes411 412 if len(annos) and "segmentation" in annos[0]:413 segms = [obj["segmentation"] for obj in annos]414 if mask_format == "polygon":415 try:416 masks = PolygonMasks(segms)417 except ValueError as e:418 raise ValueError(419 "Failed to use mask_format=='polygon' from the given annotations!"420 ) from e421 else:422 assert mask_format == "bitmask", mask_format423 masks = []424 for segm in segms:425 if isinstance(segm, list):426 # polygon427 masks.append(polygons_to_bitmask(segm, *image_size))428 elif isinstance(segm, dict):429 # COCO RLE430 masks.append(mask_util.decode(segm))431 elif isinstance(segm, np.ndarray):432 assert segm.ndim == 2, "Expect segmentation of 2 dimensions, got {}.".format(433 segm.ndim434 )435 # mask array436 masks.append(segm)437 else:438 raise ValueError(439 "Cannot convert segmentation of type '{}' to BitMasks!"440 "Supported types are: polygons as list[list[float] or ndarray],"441 " COCO-style RLE as a dict, or a binary segmentation mask "442 " in a 2D numpy array of shape HxW.".format(type(segm))443 )444 # torch.from_numpy does not support array with negative stride.445 masks = BitMasks(446 torch.stack([torch.from_numpy(np.ascontiguousarray(x)) for x in masks])447 )448 target.gt_masks = masks449 450 if len(annos) and "keypoints" in annos[0]:451 kpts = [obj.get("keypoints", []) for obj in annos]452 target.gt_keypoints = Keypoints(kpts)453 454 return target455 456 457def annotations_to_instances_rotated(annos, image_size):458 """459 Create an :class:`Instances` object used by the models,460 from instance annotations in the dataset dict.461 Compared to `annotations_to_instances`, this function is for rotated boxes only462 463 Args:464 annos (list[dict]): a list of instance annotations in one image, each465 element for one instance.466 image_size (tuple): height, width467 468 Returns:469 Instances:470 Containing fields "gt_boxes", "gt_classes",471 if they can be obtained from `annos`.472 This is the format that builtin models expect.473 """474 boxes = [obj["bbox"] for obj in annos]475 target = Instances(image_size)476 boxes = target.gt_boxes = RotatedBoxes(boxes)477 boxes.clip(image_size)478 479 classes = [obj["category_id"] for obj in annos]480 classes = torch.tensor(classes, dtype=torch.int64)481 target.gt_classes = classes482 483 return target484 485 486def filter_empty_instances(487 instances, by_box=True, by_mask=True, box_threshold=1e-5, return_mask=False488):489 """490 Filter out empty instances in an `Instances` object.491 492 Args:493 instances (Instances):494 by_box (bool): whether to filter out instances with empty boxes495 by_mask (bool): whether to filter out instances with empty masks496 box_threshold (float): minimum width and height to be considered non-empty497 return_mask (bool): whether to return boolean mask of filtered instances498 499 Returns:500 Instances: the filtered instances.501 tensor[bool], optional: boolean mask of filtered instances502 """503 assert by_box or by_mask504 r = []505 if by_box:506 r.append(instances.gt_boxes.nonempty(threshold=box_threshold))507 if instances.has("gt_masks") and by_mask:508 r.append(instances.gt_masks.nonempty())509 510 # TODO: can also filter visible keypoints511 512 if not r:513 return instances514 m = r[0]515 for x in r[1:]:516 m = m & x517 if return_mask:518 return instances[m], m519 return instances[m]520 521 522def create_keypoint_hflip_indices(dataset_names: Union[str, List[str]]) -> List[int]:523 """524 Args:525 dataset_names: list of dataset names526 527 Returns:528 list[int]: a list of size=#keypoints, storing the529 horizontally-flipped keypoint indices.530 """531 if isinstance(dataset_names, str):532 dataset_names = [dataset_names]533 534 check_metadata_consistency("keypoint_names", dataset_names)535 check_metadata_consistency("keypoint_flip_map", dataset_names)536 537 meta = MetadataCatalog.get(dataset_names[0])538 names = meta.keypoint_names539 # TODO flip -> hflip540 flip_map = dict(meta.keypoint_flip_map)541 flip_map.update({v: k for k, v in flip_map.items()})542 flipped_names = [i if i not in flip_map else flip_map[i] for i in names]543 flip_indices = [names.index(i) for i in flipped_names]544 return flip_indices545 546 547def get_fed_loss_cls_weights(dataset_names: Union[str, List[str]], freq_weight_power=1.0):548 """549 Get frequency weight for each class sorted by class id.550 We now calcualte freqency weight using image_count to the power freq_weight_power.551 552 Args:553 dataset_names: list of dataset names554 freq_weight_power: power value555 """556 if isinstance(dataset_names, str):557 dataset_names = [dataset_names]558 559 check_metadata_consistency("class_image_count", dataset_names)560 561 meta = MetadataCatalog.get(dataset_names[0])562 class_freq_meta = meta.class_image_count563 class_freq = torch.tensor(564 [c["image_count"] for c in sorted(class_freq_meta, key=lambda x: x["id"])]565 )566 class_freq_weight = class_freq.float() ** freq_weight_power567 return class_freq_weight568 569 570def gen_crop_transform_with_instance(crop_size, image_size, instance):571 """572 Generate a CropTransform so that the cropping region contains573 the center of the given instance.574 575 Args:576 crop_size (tuple): h, w in pixels577 image_size (tuple): h, w578 instance (dict): an annotation dict of one instance, in Detectron2's579 dataset format.580 """581 crop_size = np.asarray(crop_size, dtype=np.int32)582 bbox = BoxMode.convert(instance["bbox"], instance["bbox_mode"], BoxMode.XYXY_ABS)583 center_yx = (bbox[1] + bbox[3]) * 0.5, (bbox[0] + bbox[2]) * 0.5584 assert (585 image_size[0] >= center_yx[0] and image_size[1] >= center_yx[1]586 ), "The annotation bounding box is outside of the image!"587 assert (588 image_size[0] >= crop_size[0] and image_size[1] >= crop_size[1]589 ), "Crop size is larger than image size!"590 591 min_yx = np.maximum(np.floor(center_yx).astype(np.int32) - crop_size, 0)592 max_yx = np.maximum(np.asarray(image_size, dtype=np.int32) - crop_size, 0)593 max_yx = np.minimum(max_yx, np.ceil(center_yx).astype(np.int32))594 595 y0 = np.random.randint(min_yx[0], max_yx[0] + 1)596 x0 = np.random.randint(min_yx[1], max_yx[1] + 1)597 return T.CropTransform(x0, y0, crop_size[1], crop_size[0])598 599 600def check_metadata_consistency(key, dataset_names):601 """602 Check that the datasets have consistent metadata.603 604 Args:605 key (str): a metadata key606 dataset_names (list[str]): a list of dataset names607 608 Raises:609 AttributeError: if the key does not exist in the metadata610 ValueError: if the given datasets do not have the same metadata values defined by key611 """612 if len(dataset_names) == 0:613 return614 logger = logging.getLogger(__name__)615 entries_per_dataset = [getattr(MetadataCatalog.get(d), key) for d in dataset_names]616 for idx, entry in enumerate(entries_per_dataset):617 if entry != entries_per_dataset[0]:618 logger.error(619 "Metadata '{}' for dataset '{}' is '{}'".format(key, dataset_names[idx], str(entry))620 )621 logger.error(622 "Metadata '{}' for dataset '{}' is '{}'".format(623 key, dataset_names[0], str(entries_per_dataset[0])624 )625 )626 raise ValueError("Datasets have different metadata '{}'!".format(key))627 628 629def build_augmentation(cfg, is_train):630 """631 Create a list of default :class:`Augmentation` from config.632 Now it includes resizing and flipping.633 634 Returns:635 list[Augmentation]636 """637 if is_train:638 min_size = cfg.INPUT.MIN_SIZE_TRAIN639 max_size = cfg.INPUT.MAX_SIZE_TRAIN640 sample_style = cfg.INPUT.MIN_SIZE_TRAIN_SAMPLING641 else:642 min_size = cfg.INPUT.MIN_SIZE_TEST643 max_size = cfg.INPUT.MAX_SIZE_TEST644 sample_style = "choice"645 augmentation = [T.ResizeShortestEdge(min_size, max_size, sample_style)]646 if is_train and cfg.INPUT.RANDOM_FLIP != "none":647 augmentation.append(648 T.RandomFlip(649 horizontal=cfg.INPUT.RANDOM_FLIP == "horizontal",650 vertical=cfg.INPUT.RANDOM_FLIP == "vertical",651 )652 )653 return augmentation654 655 656build_transform_gen = build_augmentation657"""658Alias for backward-compatibility.659"""660