Spaces:
Build error
Build error
File size: 2,278 Bytes
1865436 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
# coding:utf-8
import os
import json
import numpy as np
import torch.utils.data as data
from detectron2.structures import (
Boxes,
PolygonMasks,
BoxMode
)
DATASETS = {
"coco_2017_train": {
"img_dir": "coco/train2017",
"ann_file": "coco/annotations/instances_train2017.json"
},
"coco_2017_val": {
"img_dir": "coco/val2017",
"ann_file": "coco/annotations/instances_val2017.json"
}
}
class MaskLoader(data.Dataset):
"""
Dataloader for Local Mask.
Arguments:
root (string): filepath to dataset folder.
dataset (string): mask to use (eg. 'train', 'val').
size (tuple): The size used for train/val (height, width).
transform (callable, optional): transformation to perform on the input mask.
"""
def __init__(self, root="datasets", dataset="coco_2017_train", size=28, transform=False):
self.root = root
self.dataset = dataset
self.transform = transform
if isinstance(size, int):
self.size = size
else:
raise TypeError
data_info = DATASETS[dataset]
img_dir, ann_file = data_info['img_dir'], data_info['ann_file']
img_dir = os.path.join(self.root, img_dir) # actually we do not use it.
ann_file = os.path.join(self.root, ann_file)
with open(ann_file, 'r') as f:
anns = json.load(f)
anns = anns['annotations']
coco = list()
for ann in anns:
if ann.get('iscrowd', 0) == 0:
coco.append(ann)
self.coco = coco
print("Removed {} images with no usable annotations. {} images left.".format(
len(anns) - len(self.coco), len(self.coco)))
def __len__(self):
return len(self.coco)
def __getitem__(self, index):
ann = self.coco[index]
# bbox transform.
bbox = np.array([ann["bbox"]]) # xmin, ymin, w, h
bbox = BoxMode.convert(bbox, BoxMode.XYWH_ABS, BoxMode.XYXY_ABS) # x1y1x2y2
bbox = Boxes(bbox)
# label
# mask transform.
mask = PolygonMasks([ann["segmentation"]])
mask = mask.crop_and_resize(bbox.tensor, self.size).float()
return mask
|