Spaces:
Sleeping
Sleeping
File size: 1,812 Bytes
2755e66 |
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 |
import os
import torch
import numpy as np
from PIL import Image
from pycocotools.coco import COCO
import torchvision.transforms as T
class LegoDataset(torch.utils.data.Dataset):
def __init__(self, root, annFile, transforms=None):
self.root = root
self.coco = COCO(annFile)
self.ids = list(self.coco.imgs.keys())
self.transforms = transforms or T.Compose([T.ToTensor()])
def __getitem__(self, index):
img_id = self.ids[index]
img_info = self.coco.loadImgs(img_id)[0]
path = img_info["file_name"]
img = Image.open(os.path.join(self.root, path)).convert("RGB")
ann_ids = self.coco.getAnnIds(imgIds=img_id)
annotations = self.coco.loadAnns(ann_ids)
boxes = []
labels = []
masks = [] # Dummy masks
for ann in annotations:
xmin, ymin, width, height = ann["bbox"]
boxes.append([xmin, ymin, xmin + width, ymin + height])
labels.append(1) # 'lego' is the only class, labeled as 1
# Dummy mask for Mask R-CNN, filled with zeros
dummy_mask = np.zeros(
(img_info["height"], img_info["width"]), dtype=np.uint8
)
masks.append(dummy_mask)
boxes = torch.as_tensor(boxes, dtype=torch.float32)
labels = torch.as_tensor(labels, dtype=torch.int64)
masks = torch.as_tensor(np.array(masks), dtype=torch.uint8)
target = {
"boxes": boxes,
"labels": labels,
"masks": masks,
"image_id": torch.tensor([img_id]),
}
if self.transforms:
img = self.transforms(img)
return img, target
def __len__(self):
return len(self.ids)
|