|
- # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
- import json
- from pathlib import Path
-
- import numpy as np
- import torch
- from PIL import Image
-
- from panopticapi.utils import rgb2id
- from util.box_ops import masks_to_boxes
-
- from .coco import make_coco_transforms
-
-
- class CocoPanoptic:
- def __init__(self, img_folder, ann_folder, ann_file, transforms=None, return_masks=True):
- with open(ann_file, 'r') as f:
- self.coco = json.load(f)
-
- # sort 'images' field so that they are aligned with 'annotations'
- # i.e., in alphabetical order
- self.coco['images'] = sorted(self.coco['images'], key=lambda x: x['id'])
- # sanity check
- if "annotations" in self.coco:
- for img, ann in zip(self.coco['images'], self.coco['annotations']):
- assert img['file_name'][:-4] == ann['file_name'][:-4]
-
- self.img_folder = img_folder
- self.ann_folder = ann_folder
- self.ann_file = ann_file
- self.transforms = transforms
- self.return_masks = return_masks
-
- def __getitem__(self, idx):
- ann_info = self.coco['annotations'][idx] if "annotations" in self.coco else self.coco['images'][idx]
- img_path = Path(self.img_folder) / ann_info['file_name'].replace('.png', '.jpg')
- ann_path = Path(self.ann_folder) / ann_info['file_name']
-
- img = Image.open(img_path).convert('RGB')
- w, h = img.size
- if "segments_info" in ann_info:
- masks = np.asarray(Image.open(ann_path), dtype=np.uint32)
- masks = rgb2id(masks)
-
- ids = np.array([ann['id'] for ann in ann_info['segments_info']])
- masks = masks == ids[:, None, None]
-
- masks = torch.as_tensor(masks, dtype=torch.uint8)
- labels = torch.tensor([ann['category_id'] for ann in ann_info['segments_info']], dtype=torch.int64)
-
- target = {}
- target['image_id'] = torch.tensor([ann_info['image_id'] if "image_id" in ann_info else ann_info["id"]])
- if self.return_masks:
- target['masks'] = masks
- target['labels'] = labels
-
- target["boxes"] = masks_to_boxes(masks)
-
- target['size'] = torch.as_tensor([int(h), int(w)])
- target['orig_size'] = torch.as_tensor([int(h), int(w)])
- if "segments_info" in ann_info:
- for name in ['iscrowd', 'area']:
- target[name] = torch.tensor([ann[name] for ann in ann_info['segments_info']])
-
- if self.transforms is not None:
- img, target = self.transforms(img, target)
-
- return img, target
-
- def __len__(self):
- return len(self.coco['images'])
-
- def get_height_and_width(self, idx):
- img_info = self.coco['images'][idx]
- height = img_info['height']
- width = img_info['width']
- return height, width
-
-
- def build(image_set, args):
- img_folder_root = Path(args.coco_path)
- ann_folder_root = Path(args.coco_panoptic_path)
- assert img_folder_root.exists(), f'provided COCO path {img_folder_root} does not exist'
- assert ann_folder_root.exists(), f'provided COCO path {ann_folder_root} does not exist'
- mode = 'panoptic'
- PATHS = {
- "train": ("train2017", Path("annotations") / f'{mode}_train2017.json'),
- "val": ("val2017", Path("annotations") / f'{mode}_val2017.json'),
- }
-
- img_folder, ann_file = PATHS[image_set]
- img_folder_path = img_folder_root / img_folder
- ann_folder = ann_folder_root / f'{mode}_{img_folder}'
- ann_file = ann_folder_root / ann_file
-
- dataset = CocoPanoptic(img_folder_path, ann_folder, ann_file,
- transforms=make_coco_transforms(image_set), return_masks=args.masks)
-
- return dataset
|