diff --git a/Streamlit/section.py b/Streamlit/section.py index 6371827..14b99be 100644 --- a/Streamlit/section.py +++ b/Streamlit/section.py @@ -30,7 +30,6 @@ PASCAL_VOC_FOLDER_IMAGES = 'dataset/sample_format_pascal_voc/JPEGImages/' PASCAL_VOC_FOLDER_ANNOTATIONS = 'dataset/sample_format_pascal_voc/Annotations/' PASCAL_VOC_FILE_IMAGESETS = 'dataset/sample_format_pascal_voc/ImageSets/Main/default.txt' -PASCAL_VOC_FILE_CLASSES = 'dataset/sample_format_pascal_voc/classes.txt' DATASET_SAMPLE = 'Sample' DATASET_FORMAT_YOLO = 'YOLO' @@ -110,8 +109,7 @@ def dataset_pascal_voc(key: str) -> Dataset: folder_images = st.text_input('Path to Pascal VOC Images', PASCAL_VOC_FOLDER_IMAGES, key=f'{key}-pascal-voc-folder_images') folder_annotations = st.text_input('Path to Pascal VOC Annotations', PASCAL_VOC_FOLDER_ANNOTATIONS, key=f'{key}-pascal-voc-folder_annotations') file_imagesets = st.text_input('Path to Pascal VOC Image Sets', PASCAL_VOC_FILE_IMAGESETS, key=f'{key}-pascal-voc-file_imagesets') - file_classes = st.text_input('Path to Pascal VOC Classes', PASCAL_VOC_FILE_CLASSES, key=f'{key}-pascal-voc-file_classes') - return loader.load_pascal_voc_dataset(folder_images, folder_annotations, file_imagesets, file_classes) + return loader.load_pascal_voc_dataset(folder_images, folder_annotations, file_imagesets) def background_image(key: str) -> Image: st.subheader('Background Image') diff --git a/patchmentation/utils/loader.py b/patchmentation/utils/loader.py index bb71484..afc8fc5 100644 --- a/patchmentation/utils/loader.py +++ b/patchmentation/utils/loader.py @@ -5,6 +5,7 @@ import numpy as np import tempfile import json +import xml.etree.ElementTree as ET from typing import Dict, List, Union, Any, Tuple temporary_folder = tempfile.TemporaryDirectory() @@ -175,5 +176,87 @@ def convert_coco_bbox(x: int, y: int, width: int, height: int) -> BBox: bbox = BBox(xmin, ymin, xmax, ymax) return bbox -def load_pascal_voc_dataset(folder_images: str, folder_annotations: str, file_imagesets: str, file_classes: str) -> Dataset: - pass +def load_pascal_voc_dataset(folder_images: str, folder_annotations: str, file_imagesets: str) -> Dataset: + imagesets = load_pascal_voc_imagesets(file_imagesets) + image_patches = load_pascal_voc_image_patches(folder_images, folder_annotations, imagesets) + classes = set() + for _, patches in image_patches: + for _, _, class_name in patches: + classes.add(class_name) + classes = list(classes) + dataset = Dataset(image_patches, classes) + return dataset + +def load_pascal_voc_imagesets(file_imagesets: str) -> List[str]: + imagesets = [] + with open(file_imagesets, 'r') as f: + lines = f.readlines() + for line in lines: + line = line.strip() + imagesets.append(line) + return imagesets + +def load_pascal_voc_image_patches(folder_images: str, folder_annotations: str, imagesets: List[str]) -> List[ImagePatch]: + images = load_pacal_voc_images(folder_images, imagesets) + annotations = load_pascal_voc_annotations(folder_annotations, imagesets) + image_patches = [] + for image, annotation in zip(images, annotations): + patches = [] + for bbox, class_name in annotation: + patch = Patch(image, bbox, class_name) + patches.append(patch) + image_patch = ImagePatch(image, patches) + image_patches.append(image_patch) + return image_patches + +def load_pacal_voc_images(folder_images: str, imagesets: List[str]) -> List[Image]: + images = [] + for file_name in imagesets: + path = None + for ext in ['.jpg', '.png', '.jpeg', '']: + path = os.path.join(folder_images, file_name + ext) + if os.path.exists(path): + break + assert os.path.exists(path) + image = Image(path) + images.append(image) + return images + +def load_pascal_voc_annotations(folder_annotations: str, imagesets: List[str]) -> List[List[Tuple[BBox, str]]]: + annotations = [] + for file_name in imagesets: + path = None + ext = '.xml' + path = os.path.join(folder_annotations, file_name + ext) + assert os.path.exists(path) + annotation = load_pascal_voc_xml(path) + annotations.append(annotation) + return annotations + +def load_pascal_voc_xml(file_xml: str) -> List[Tuple[BBox, str]]: + annotation = [] + tree = ET.parse(file_xml) + root = tree.getroot() + for child in root: + if child.tag == 'object': + class_name = xmin = ymin = xmax = ymax = None + for attribute in child: + if attribute.tag == 'name': + class_name = attribute.text + if attribute.tag == 'bndbox': + for cord in attribute: + if cord.tag == 'xmin': xmin = int(float(cord.text)) + if cord.tag == 'ymin': ymin = int(float(cord.text)) + if cord.tag == 'xmax': xmax = int(float(cord.text)) + if cord.tag == 'ymax': ymax = int(float(cord.text)) + bbox = convert_pascal_voc_bbox(xmin, ymin, xmax, ymax) + annotation.append((bbox, class_name)) + return annotation + +def convert_pascal_voc_bbox(xmin: int, ymin: int, xmax: int, ymax: int) -> BBox: + xmin = int(xmin) + ymin = int(ymin) + xmax = int(xmax) + ymax = int(ymax) + bbox = BBox(xmin, ymin, xmax, ymax) + return bbox diff --git a/tests/test_utils/test_loader.py b/tests/test_utils/test_loader.py index 2fb2c51..192bf99 100644 --- a/tests/test_utils/test_loader.py +++ b/tests/test_utils/test_loader.py @@ -6,9 +6,6 @@ from patchmentation.utils import validator from patchmentation.collections import BBox -import pytest -import numpy as np - YOLO_FOLDER_IMAGES = 'dataset/sample_format_yolo/obj_train_data/' YOLO_FOLDER_ANNOTATIONS = 'dataset/sample_format_yolo/obj_train_data/' YOLO_FILE_NAMES = 'dataset/sample_format_yolo/obj.names' @@ -19,7 +16,6 @@ PASCAL_VOC_FOLDER_IMAGES = 'dataset/sample_format_pascal_voc/JPEGImages/' PASCAL_VOC_FOLDER_ANNOTATIONS = 'dataset/sample_format_pascal_voc/Annotations/' PASCAL_VOC_FILE_IMAGESETS = 'dataset/sample_format_pascal_voc/ImageSets/Main/default.txt' -PASCAL_VOC_FILE_CLASSES = 'dataset/sample_format_pascal_voc/classes.txt' def test_save_load_image_array_1(): image_array = helper.generate_image_array() @@ -65,7 +61,6 @@ def test_loader_coco(): dataset = loader.load_coco_dataset(COCO_FOLDER_IMAGES, COCO_FILE_ANNOTATIONS) validator.validate_Dataset(dataset, check_image_bbox=True) -@pytest.mark.skip('Not implemented') def test_loader_pascal_voc(): - dataset = loader.load_pascal_voc_dataset(PASCAL_VOC_FOLDER_IMAGES, PASCAL_VOC_FOLDER_ANNOTATIONS, PASCAL_VOC_FILE_IMAGESETS, PASCAL_VOC_FILE_CLASSES) + dataset = loader.load_pascal_voc_dataset(PASCAL_VOC_FOLDER_IMAGES, PASCAL_VOC_FOLDER_ANNOTATIONS, PASCAL_VOC_FILE_IMAGESETS) validator.validate_Dataset(dataset, check_image_bbox=True)