Spaces:
Sleeping
Sleeping
File size: 1,633 Bytes
749745d |
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 |
import os
import os.path
import json
from PIL import Image
import torch
import torchvision
import torch.utils.data as data
from maskrcnn_benchmark.structures.bounding_box import BoxList
class Background(data.Dataset):
"""Background
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``
"""
def __init__(self, ann_file, root, remove_images_without_annotations=None, transforms=None):
self.root = root
with open(ann_file, "r") as f:
self.ids = json.load(f)["images"]
self.transform = transforms
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
tuple: Tuple (image, target). target is the object returned by ``coco.loadAnns``.
"""
im_info = self.ids[index]
path = im_info["file_name"]
fp = os.path.join(self.root, path)
img = Image.open(fp).convert("RGB")
if self.transform is not None:
img, _ = self.transform(img, None)
null_target = BoxList(torch.zeros((0, 4)), (img.shape[-1], img.shape[-2]))
null_target.add_field("labels", torch.zeros(0))
return img, null_target, index
def __len__(self):
return len(self.ids)
def get_img_info(self, index):
im_info = self.ids[index]
return im_info
|