Added src files

6ca1e94 verified 3 months ago

7.32 kB

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)

	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)