File size: 7,320 Bytes
6ca1e94 | 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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 | 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) |