class CocoDetection(Dataset): def __init__(self, root, ann_file, transform=None): from hpycocotools.coco import COCO self.root = root self.ann_file = ann_file with open(self.ann_file, 'r') as f: self.data = json.load(f) self.coco = COCO(self.data, verbose=False) self.ids = list(self.coco.imgs.keys()) self.transform = transform def to_coco(self, indices=None): if indices is None: return self.data ids = [self.ids[i] for i in indices] images = self.coco.loadImgs(ids) ann_ids = self.coco.getAnnIds(ids) annotations = self.coco.loadAnns(ann_ids) return { **self.data, "images": images, "annotations": annotations, } def __getitem__(self, index): """ 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] ann_ids = coco.getAnnIds(imgIds=img_id) target = coco.loadAnns(ann_ids) path = coco.loadImgs(img_id)[0]['file_name'] img = Image.open(os.path.join(self.root, path)).convert('RGB') if self.transform is not None: img, target = self.transform(img, target) 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))) return fmt_str
class COCOEval(Metric): def __init__(self, annotations, iou_type='bbox', percent_area_ranges=None): self.annotations = annotations self.iou_type = iou_type self.percent_area_ranges = percent_area_ranges super().__init__() def reset(self): from hpycocotools.coco import COCO self.coco_gt = COCO(self.annotations, verbose=False, use_percent=bool(self.percent_area_ranges)) self.res = [] def update(self, output): target, image_dets, batch_size = get(["target", "preds", "batch_size"], output) image_gts = target[0] for dets, gts in zip(image_dets, image_gts): image_id = gts[0]['image_id'] for d in dets: d = {**d, 'image_id': image_id} self.res.append(d) def compute(self): from hpycocotools.coco import COCO img_ids = list(set([d['image_id'] for d in self.res])) imgs = self.coco_gt.loadImgs(img_ids) ann_ids = self.coco_gt.getAnnIds(imgIds=img_ids) anns = self.coco_gt.loadAnns(ann_ids) annotations = { **self.annotations, 'images': imgs, 'annotations': anns, } coco_gt = COCO(annotations, verbose=False, use_percent=bool(self.percent_area_ranges)) from hpycocotools.cocoeval import COCOeval from hpycocotools.mask import encode for dt in self.res: img = self.coco_gt.imgs[dt['image_id']] width = img['width'] height = img['height'] l, t, w, h = dt['bbox'] l *= width t *= height w *= width h *= height dt['bbox'] = [l, t, w, h] if 'segmentation' in dt and self.iou_type == 'segm': r = int(l + w) b = int(t + h) l = max(0, int(l)) t = max(0, int(t)) r = min(r, width) b = min(b, height) w = r - l h = b - t m = np.zeros((height, width), dtype=np.uint8) segm = cv2.resize(dt['segmentation'], (w, h), interpolation=cv2.INTER_NEAREST) m[t:b, l:r] = segm dt['segmentation'] = encode(np.asfortranarray(m)) coco_dt = coco_gt.loadRes(self.res) ev = COCOeval(coco_gt, coco_dt, iouType=self.iou_type, pAreaRng=self.percent_area_ranges, verbose=False) ev.evaluate() ev.accumulate() ev.summarize() return ev.stats[0]
class VOCDetection(Dataset): """`Pascal VOC <http://host.robots.ox.ac.uk/pascal/VOC/>`_ Detection Dataset. Args: root (string): Root directory of the VOC Dataset. year (string, optional): The dataset year, supports years 2007 to 2012. image_set (string, optional): Select the image_set to use, ``trainval`` or ``test`` download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. (default: alphabetic indexing of VOC's 20 classes). transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` """ def __init__(self, root, year='2012', image_set='trainval', download=False, transform=None): self.root = Path(root).expanduser().absolute() self.year = year self.image_set = image_set if image_set == 'test': dataset_dict = TEST_DATASET_YEAR_DICT else: dataset_dict = DATASET_YEAR_DICT self.url = dataset_dict[year]['url'] self.filename = dataset_dict[year]['filename'] self.md5 = dataset_dict[year]['md5'] base_dir = dataset_dict[year]['base_dir'] self.voc_root = self.root / base_dir self.image_dir = self.voc_root / 'JPEGImages' self.ann_file_url = dataset_dict[year]['ann_file_url'] self.ann_file = self.voc_root / ("%s%s.json" % (image_set, year)) if download: self.download() from hpycocotools.coco import COCO with open(self.ann_file, 'r') as f: self.data = json.load(f) self.coco = COCO(self.data, verbose=False) self.ids = list(self.coco.imgs.keys()) self.transform = transform def to_coco(self, indices=None): if indices is None: return self.data ids = [self.ids[i] for i in indices] images = self.coco.loadImgs(ids) ann_ids = self.coco.getAnnIds(ids) annotations = self.coco.loadAnns(ann_ids) return { **self.data, "images": images, "annotations": annotations, } def __getitem__(self, index): """ 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] ann_ids = coco.getAnnIds(imgIds=img_id) anns = coco.loadAnns(ann_ids) path = coco.loadImgs([img_id])[0]['file_name'] img = Image.open(self.image_dir / path).convert('RGB') if self.transform is not None: img, anns = self.transform(img, anns) return img, anns def __len__(self): return len(self.ids) def download(self): import tarfile if self.voc_root.is_dir() and self.ann_file.exists(): print("Dataset found. Skip download or extract.") return if not self.voc_root.is_dir(): download_url(self.url, self.root, self.filename, self.md5) with tarfile.open(self.root / self.filename, "r") as tar: tar.extractall(path=self.root) if not self.ann_file.exists(): google_drive_match = re.match( r"https://drive.google.com/open\?id=(.*)", self.ann_file_url) file_id = google_drive_match.group(1) download_google_drive(file_id, self.voc_root, self.ann_file.name) def __repr__(self): fmt_str = 'Dataset ' + self.__class__.__name__ + '\n' fmt_str += ' Year: {}\n'.format(self.year) fmt_str += ' ImageSet: {}\n'.format(self.image_set) 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))) return fmt_str
class SVHNDetection(Dataset): """`SVHN <http://ufldl.stanford.edu/housenumbers/>`_ Dataset. Note: The SVHN dataset assigns the label `10` to the digit `0`. However, in this Dataset, we assign the label `0` to the digit `0` to be compatible with PyTorch loss functions which expect the class labels to be in the range `[0, C-1]` Args: root (string): Root directory of dataset where directory ``SVHN`` exists. split (string): One of {'train', 'test', 'extra'}. Accordingly dataset is selected. 'extra' is Extra training set. transform (callable, optional): A function/transform that takes in an PIL image and targets simultaneously and returns a transformed version of them. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ def __init__(self, root, split='train', transform=None, download=False): self.root = Path(root).expanduser().absolute() self.split = split self.transform = transform self.filename = SPLIT_FILES[split]['name'] self.md5 = SPLIT_FILES[split]['md5'] self.url = SPLIT_FILES[split]["url"] self.img_dir = self.root / split self.ann_filename = split + ".json" self.ann_dir = self.root / "annotations" self.ann_dir.mkdir(exist_ok=True, parents=True) self.ann_file = self.ann_dir / self.ann_filename if download: self.download() with open(self.ann_file, 'r') as f: self.data = json.load(f) from hpycocotools.coco import COCO self.coco = COCO(self.data, verbose=False) self.ids = list(self.coco.imgs.keys()) def to_coco(self, indices=None): if indices is None: return self.data ids = [self.ids[i] for i in indices] images = self.coco.loadImgs(ids) ann_ids = self.coco.getAnnIds(ids) annotations = self.coco.loadAnns(ann_ids) return { **self.data, "images": images, "annotations": annotations, } def __getitem__(self, index): """ Args: index (int): Index Returns: tuple: (image, anns) where target is a dictionary of the XML tree. """ coco = self.coco img_id = self.ids[index] ann_ids = coco.getAnnIds(imgIds=img_id) target = coco.loadAnns(ann_ids) path = coco.loadImgs(img_id)[0]['file_name'] img = Image.open(self.img_dir / path).convert('RGB') if self.transform is not None: img, target = self.transform(img, target) return img, target def __len__(self): return len(self.ids) def download(self): if self.img_dir.is_dir() and self.ann_file.exists(): print("Dataset found. Skip download or extract") return google_drive_match = re.match( r"https://drive.google.com/open\?id=(.*)", self.url) if google_drive_match: file_id = google_drive_match.group(1) download_google_drive(file_id, self.root, self.filename, self.md5) else: download_url(self.url, self.root, self.filename, self.md5) file_path = self.root / self.filename with tarfile.open(file_path, "r") as tar: tar.extractall(path=self.root) ann_file = self.root / self.ann_filename ann_file.rename(self.ann_file)