Merge pull request #23 from Xu-Justin/feature/pascal-voc-loader

add pascal voc loader
Xu-Justin · Nov 24, 2022 · 09ea0c2 · 09ea0c2
2 parents 4ecec5a + b47ec73
commit 09ea0c2
Show file tree

Hide file tree

Showing 3 changed files with 87 additions and 11 deletions.
diff --git 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
@@ -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
@@ -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)