# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. import os import os.path import math from PIL import Image, ImageDraw import random import numpy as np import torch import torchvision import torch.utils.data as data from maskrcnn_benchmark.structures.bounding_box import BoxList from maskrcnn_benchmark.structures.segmentation_mask import SegmentationMask from maskrcnn_benchmark.structures.keypoint import PersonKeypoints from maskrcnn_benchmark.config import cfg import pdb def _count_visible_keypoints(anno): return sum(sum(1 for v in ann["keypoints"][2::3] if v > 0) for ann in anno) def _has_only_empty_bbox(anno): return all(any(o <= 1 for o in obj["bbox"][2:]) for obj in anno) def has_valid_annotation(anno): # if it's empty, there is no annotation if len(anno) == 0: return False # if all boxes have close to zero area, there is no annotation if _has_only_empty_bbox(anno): return False # keypoints task have a slight different critera for considering # if an annotation is valid if "keypoints" not in anno[0]: return True # for keypoint detection tasks, only consider valid images those # containing at least min_keypoints_per_image if _count_visible_keypoints(anno) >= cfg.DATALOADER.MIN_KPS_PER_IMS: return True return False def pil_loader(path, retry=5): # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835) ri = 0 while ri < retry: try: with open(path, "rb") as f: img = Image.open(f) return img.convert("RGB") except: ri += 1 def rgb2id(color): if isinstance(color, np.ndarray) and len(color.shape) == 3: if color.dtype == np.uint8: color = color.astype(np.int32) return color[:, :, 0] + 256 * color[:, :, 1] + 256 * 256 * color[:, :, 2] return int(color[0] + 256 * color[1] + 256 * 256 * color[2]) class CocoDetection(data.Dataset): """`MS Coco Detection `_ Dataset. Args: root (string): Root directory where images are downloaded to. annFile (string): Path to json annotation file. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.ToTensor`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. """ def __init__(self, root, annFile, transform=None, target_transform=None): from pycocotools.coco import COCO self.root = root self.coco = COCO(annFile) self.ids = list(self.coco.imgs.keys()) self.transform = transform self.target_transform = target_transform def __getitem__(self, index, return_meta=False): """ Args: index (int): Index Returns: tuple: Tuple (image, target). target is the object returned by ``coco.loadAnns``. """ coco = self.coco img_id = self.ids[index] if isinstance(img_id, str): img_id = [img_id] ann_ids = coco.getAnnIds(imgIds=img_id) target = coco.loadAnns(ann_ids) meta = coco.loadImgs(img_id)[0] path = meta["file_name"] img = pil_loader(os.path.join(self.root, path)) if self.transform is not None: img = self.transform(img) if self.target_transform is not None: target = self.target_transform(target) if return_meta: return img, target, meta else: return img, target def __len__(self): return len(self.ids) def __repr__(self): fmt_str = "Dataset " + self.__class__.__name__ + "\n" fmt_str += " Number of datapoints: {}\n".format(self.__len__()) fmt_str += " Root Location: {}\n".format(self.root) tmp = " Transforms (if any): " fmt_str += "{0}{1}\n".format(tmp, self.transform.__repr__().replace("\n", "\n" + " " * len(tmp))) tmp = " Target Transforms (if any): " fmt_str += "{0}{1}".format(tmp, self.target_transform.__repr__().replace("\n", "\n" + " " * len(tmp))) return fmt_str class COCODataset(CocoDetection): def __init__( self, ann_file, root, remove_images_without_annotations, transforms=None, ignore_crowd=True, max_box=-1, few_shot=0, one_hot=False, override_category=None, **kwargs ): super(COCODataset, self).__init__(root, ann_file) # sort indices for reproducible results self.ids = sorted(self.ids) # filter images without detection annotations if remove_images_without_annotations: ids = [] for img_id in self.ids: if isinstance(img_id, str): ann_ids = self.coco.getAnnIds(imgIds=[img_id], iscrowd=None) else: ann_ids = self.coco.getAnnIds(imgIds=img_id, iscrowd=None) anno = self.coco.loadAnns(ann_ids) if has_valid_annotation(anno): ids.append(img_id) self.ids = ids if few_shot: ids = [] cats_freq = [few_shot] * len(self.coco.cats.keys()) if "shuffle_seed" in kwargs and kwargs["shuffle_seed"] != 0: import random random.Random(kwargs["shuffle_seed"]).shuffle(self.ids) print("Shuffle the dataset with random seed: ", kwargs["shuffle_seed"]) for img_id in self.ids: if isinstance(img_id, str): ann_ids = self.coco.getAnnIds(imgIds=[img_id], iscrowd=None) else: ann_ids = self.coco.getAnnIds(imgIds=img_id, iscrowd=None) anno = self.coco.loadAnns(ann_ids) cat = set([ann["category_id"] for ann in anno]) # set/tuple corresponde to instance/image level is_needed = sum([cats_freq[c - 1] > 0 for c in cat]) if is_needed: ids.append(img_id) for c in cat: cats_freq[c - 1] -= 1 # print(cat, cats_freq) self.ids = ids if override_category is not None: self.coco.dataset["categories"] = override_category print("Override category: ", override_category) self.json_category_id_to_contiguous_id = {v: i + 1 for i, v in enumerate(self.coco.getCatIds())} self.contiguous_category_id_to_json_id = {v: k for k, v in self.json_category_id_to_contiguous_id.items()} self.id_to_img_map = {k: v for k, v in enumerate(self.ids)} self.transforms = transforms self.ignore_crowd = ignore_crowd self.max_box = max_box self.one_hot = one_hot def categories(self, no_background=True): categories = self.coco.dataset["categories"] label_list = {} for index, i in enumerate(categories): if not no_background or (i["name"] != "__background__" and i["id"] != 0): label_list[self.json_category_id_to_contiguous_id[i["id"]]] = i["name"] return label_list def __getitem__(self, idx): img, anno = super(COCODataset, self).__getitem__(idx) # filter crowd annotations if self.ignore_crowd: anno = [obj for obj in anno if obj["iscrowd"] == 0] boxes = [obj["bbox"] for obj in anno] boxes = torch.as_tensor(boxes).reshape(-1, 4) # guard against no boxes if self.max_box > 0 and len(boxes) > self.max_box: rand_idx = torch.randperm(self.max_box) boxes = boxes[rand_idx, :] else: rand_idx = None target = BoxList(boxes, img.size, mode="xywh").convert("xyxy") classes = [obj["category_id"] for obj in anno] classes = [self.json_category_id_to_contiguous_id[c] for c in classes] classes = torch.tensor(classes) if rand_idx is not None: classes = classes[rand_idx] if cfg.DATASETS.CLASS_AGNOSTIC: classes = torch.ones_like(classes) target.add_field("labels", classes) if anno and "segmentation" in anno[0]: masks = [obj["segmentation"] for obj in anno] masks = SegmentationMask(masks, img.size, mode="poly") target.add_field("masks", masks) if anno and "cbox" in anno[0]: cboxes = [obj["cbox"] for obj in anno] cboxes = torch.as_tensor(cboxes).reshape(-1, 4) # guard against no boxes cboxes = BoxList(cboxes, img.size, mode="xywh").convert("xyxy") target.add_field("cbox", cboxes) if anno and "keypoints" in anno[0]: keypoints = [] gt_keypoint = self.coco.cats[1]["keypoints"] # a better way to get keypoint description use_keypoint = cfg.MODEL.ROI_KEYPOINT_HEAD.KEYPOINT_NAME for obj in anno: if len(use_keypoint) > 0: kps = [] for name in use_keypoint: kp_idx = slice(3 * gt_keypoint.index(name), 3 * gt_keypoint.index(name) + 3) kps += obj["keypoints"][kp_idx] keypoints.append(kps) else: keypoints.append(obj["keypoints"]) keypoints = PersonKeypoints(keypoints, img.size) target.add_field("keypoints", keypoints) target = target.clip_to_image(remove_empty=True) if self.transforms is not None: img, target = self.transforms(img, target) if cfg.DATASETS.SAMPLE_RATIO != 0.0: ratio = cfg.DATASETS.SAMPLE_RATIO num_sample_target = math.ceil(len(target) * ratio) if ratio > 0 else math.ceil(-ratio) sample_idx = torch.randperm(len(target))[:num_sample_target] target = target[sample_idx] return img, target, idx def get_img_info(self, index): img_id = self.id_to_img_map[index] img_data = self.coco.imgs[img_id] return img_data