| |
| |
| |
| |
| |
| |
| |
| |
|
|
| 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) |
|
|
| |
| |
| self.coco['images'] = sorted(self.coco['images'], key=lambda x: x['id']) |
| |
| 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 |
|
|
