File size: 1,834 Bytes
b338c44 | 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 torchvision.transforms as transforms
from torch.utils.data import Dataset
from pycocotools.coco import COCO
from PIL import Image
class COCODataset(Dataset):
def __init__(self, root, annotation_file, transforms=None):
self.root = root
self.coco = COCO(annotation_file)
self.transforms = transforms
self.image_ids = list(self.coco.imgs.keys())
def __getitem__(self, idx):
image_id = self.image_ids[idx]
image_info = self.coco.imgs[image_id]
image_path = os.path.join(self.root, image_info["file_name"])
image = Image.open(image_path).convert("RGB")
annotations = self.coco.loadAnns(self.coco.getAnnIds(imgIds=image_id))
boxes = []
labels = []
for ann in annotations:
xmin = ann["bbox"][0]
ymin = ann["bbox"][1]
width = ann["bbox"][2]
height = ann["bbox"][3]
boxes.append([xmin, ymin, xmin + width, ymin + height])
labels.append(ann["category_id"]) # Ensure category_id starts from 1
if len(boxes) == 0:
boxes = torch.zeros((0, 4), dtype=torch.float32)
labels = torch.zeros((0,), dtype=torch.int64)
else:
boxes = torch.tensor(boxes, dtype=torch.float32)
labels = torch.tensor(labels, dtype=torch.int64)
target = {"boxes": boxes, "labels": labels, "image_id": torch.tensor([image_id])}
if self.transforms:
image = self.transforms(image)
return image, target
def __len__(self):
return len(self.image_ids)
def get_transforms():
return transforms.Compose([
transforms.ToTensor()
])
|