import torch import zipfile import xml.etree.ElementTree as ET import torch.nn as nn from PIL import Image from pathlib import Path from functools import partial from torch.utils.data import Dataset # class to idx in VOC dataset VOC_CLASSES = [ "aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor" ] CLS_TO_IDX = {cls_name: idx for idx, cls_name in enumerate(VOC_CLASSES)} class VOCDataset(Dataset): def __init__(self, img_paths_list, transform=None): self.img_paths_list = img_paths_list self.transform = transform def __len__(self): return len(self.img_paths_list) def __getitem__(self, idx): img_path = self.img_paths_list[idx] annotation_path = img_path.parents[1] / "Annotations" / (img_path.stem + ".xml") img = Image.open(img_path) annotation_dict = voc_to_dict(annotation_path) if self.transform: img, annotation_dict = self.transform(img, annotation_dict) return img, annotation_dict def collate_fn(batch, difficult=False): imgs, annotations = zip(*batch) max_height = max(img.shape[1] for img in imgs) max_width = max(img.shape[2] for img in imgs) padded_imgs = [] img_sizes_before_pad = [] gt_boxes = [] gt_labels = [] gt_difficult = [] # only populated when difficult=True for img, annotation in zip(imgs, annotations): img_height, img_width = img.shape[1], img.shape[2] pad_height = max_height - img_height pad_width = max_width - img_width img_sizes_before_pad.append((img_height, img_width)) # Store original height and width padded_img = nn.functional.pad(img, (0, pad_width, 0, pad_height), mode='constant', value=0) padded_imgs.append(padded_img) boxes = torch.tensor([ [obj["bndbox"]["x_min"]* img_width, obj["bndbox"]["y_min"]* img_height, obj["bndbox"]["x_max"]* img_width, obj["bndbox"]["y_max"]* img_height] for obj in annotation["objects"]], dtype=torch.float32) labels = torch.tensor([obj["class_idx"] for obj in annotation["objects"]], dtype=torch.int64) gt_boxes.append(boxes) gt_labels.append(labels) if difficult: # index-parallel to boxes/labels for this image, same ragged per-image convention gt_difficult.append(torch.tensor([obj["difficult"] for obj in annotation["objects"]], dtype=torch.bool)) if difficult: return torch.stack(padded_imgs, dim=0), gt_boxes, gt_labels, img_sizes_before_pad, gt_difficult return torch.stack(padded_imgs, dim=0), gt_boxes, gt_labels, img_sizes_before_pad def create_voc_dataloader(img_paths_list, transform=None, batch_size=32, shuffle=True, difficult=False): voc_dataset = VOCDataset(img_paths_list, transform=transform) # partial (not a lambda) so the collate stays picklable if num_workers > 0 is ever used dataloader = torch.utils.data.DataLoader(voc_dataset, batch_size=batch_size, shuffle=shuffle, collate_fn=partial(collate_fn, difficult=difficult)) return dataloader # Converting VOC annotation XML to dictionary format def voc_to_dict(annotation_path): tree = ET.parse(annotation_path) root = tree.getroot() img_width = float(root.find("size").find("width").text) img_height = float(root.find("size").find("height").text) annotation_data = { "filename" : root.find("filename").text, "size" : { "width" : img_width, "height" : img_height, "depth" : int(root.find("size").find("depth").text) }, "objects" : [] } for obj in root.findall("object"): name = obj.find("name").text difficult_tag = obj.find("difficult") # absent in a few VOC XMLs, treat as not difficult obj_dict = { "name" : name, "class_idx" : CLS_TO_IDX[name], "difficult" : int(difficult_tag.text) if difficult_tag is not None else 0, "bndbox" : { "x_min" : float(obj.find("bndbox").find("xmin").text)/img_width, "y_min" : float(obj.find("bndbox").find("ymin").text)/img_height, "x_max" : float(obj.find("bndbox").find("xmax").text)/img_width, "y_max" : float(obj.find("bndbox").find("ymax").text)/img_height } } annotation_data["objects"].append(obj_dict) return annotation_data def get_voc_img_paths_train(): data_path = Path("data/") if not data_path.exists(): raise RuntimeError("Data directory not found. Please run data_setup.py to extract the datasets.") # Define the paths to the VOC2007 and VOC2012 datasets and their image directories voc2007_path = data_path / "VOC2007" voc2012_path = data_path / "VOC2012" voc2007_img_path = voc2007_path / "JPEGImages" voc2012_img_path = voc2012_path / "JPEGImages" voc2007_trainval = voc2007_path / "ImageSets" / "Main" / "trainval.txt" voc2012_trainval = voc2012_path / "ImageSets" / "Main" / "trainval.txt" with open(voc2007_trainval, "r") as f: voc2007_img_paths_train = [voc2007_img_path / (line.strip() + ".jpg") for line in f.readlines()] with open(voc2012_trainval, "r") as f: voc2012_img_paths_train = [voc2012_img_path / ((line.strip()).split(" ")[0] + ".jpg") for line in f.readlines()] return voc2007_img_paths_train, voc2012_img_paths_train def get_voc_img_paths_test(): data_path = Path("data/") if not data_path.exists(): raise RuntimeError("Data directory not found. Please run data_setup.py to extract the datasets.") # Define the paths to the VOC2007 and VOC2012 datasets and their image directories voc2007_path = data_path / "VOC2007" voc2012_path = data_path / "VOC2012" voc2007_img_path = voc2007_path / "JPEGImages" voc2012_img_path = voc2012_path / "JPEGImages" voc2007_test = voc2007_path / "ImageSets" / "Main" / "test.txt" voc2012_test = voc2012_path / "ImageSets" / "Main" / "test.txt" with open(voc2007_test, "r") as f: voc2007_img_paths_test = [voc2007_img_path / (line.strip() + ".jpg") for line in f.readlines()] with open(voc2012_test, "r") as f: voc2012_img_paths_test = [voc2012_img_path / ((line.strip()).split(" ")[0] + ".jpg") for line in f.readlines()] return voc2007_img_paths_test, voc2012_img_paths_test if __name__ == "__main__": data_path = Path("data/") # Get the paths to the VOC2007 and VOC2012 zip files voc2007_zip_path = Path("VOC2007.zip") voc2012_zip_path = Path("VOC2012.zip") # Check if the zip files exist, if not raise an error if not voc2007_zip_path.exists() or not voc2012_zip_path.exists(): raise RuntimeError("Dataset not found.") # Extract the datasets in the data directory print("Extracting 2007 dataset ...") with zipfile.ZipFile(voc2007_zip_path, "r") as zip_ref: zip_ref.extractall(data_path) print("Extracting 2012 dataset ...") with zipfile.ZipFile(voc2012_zip_path, "r") as zip_ref: zip_ref.extractall(data_path)