Source code for omrdatasettools.AudiverisOmrImageGenerator

import argparse
import os
from glob import glob
from xml.etree import ElementTree

from PIL import Image

from omrdatasettools.Point2D import Point2D
from omrdatasettools.ExportPath import ExportPath
from omrdatasettools.Rectangle import Rectangle


class AudiverisOmrSymbol(Rectangle):
    def __init__(self, symbol_class: str, x: int, y: int, width: int, height: int) -> None:
        super().__init__(Point2D(x, y), width, height)
        self.symbol_class = symbol_class




class AudiverisOmrImageGenerator:
    def __init__(self) -> None:
        super().__init__()



    def extract_symbols(self, raw_data_directory: str, destination_directory: str):
        """
        Extracts the symbols from the raw XML documents and matching images of the Audiveris OMR dataset into
        individual symbols

        :param raw_data_directory: The directory, that contains the xml-files and matching images
        :param destination_directory: The directory, in which the symbols should be generated into. One sub-folder per
                                      symbol category will be generated automatically
        """
        print("Extracting Symbols from Audiveris OMR Dataset...")

        all_xml_files = [y for x in os.walk(raw_data_directory) for y in glob(os.path.join(x[0], '*.xml'))]
        all_image_files = [y for x in os.walk(raw_data_directory) for y in glob(os.path.join(x[0], '*.png'))]

        data_pairs = []
        for i in range(len(all_xml_files)):
            data_pairs.append((all_xml_files[i], all_image_files[i]))

        for data_pair in data_pairs:
            self.__extract_symbols(data_pair[0], data_pair[1], destination_directory)


    def __extract_symbols(self, xml_file: str, image_file: str, destination_directory: str):
        # xml_file, image_file = 'data/audiveris_omr_raw\\IMSLP06053p1.xml', 'data/audiveris_omr_raw\\IMSLP06053p1.png'
        # xml_file, image_file = 'data/audiveris_omr_raw\\mops-1.xml', 'data/audiveris_omr_raw\\mops-1.png'
        # xml_file, image_file = 'data/audiveris_omr_raw\\mtest1-1.xml', 'data/audiveris_omr_raw\\mtest1-1.png'
        # xml_file, image_file = 'data/audiveris_omr_raw\\mtest2-1.xml', 'data/audiveris_omr_raw\\mtest2-1.png'
        image = Image.open(image_file)
        annotations = ElementTree.parse(xml_file).getroot()
        xml_symbols = annotations.findall("Symbol")

        file_name_without_extension = os.path.splitext(os.path.basename(xml_file))[0]
        symbols = []

        for xml_symbol in xml_symbols:
            symbol_class = xml_symbol.get("shape")

            bounds = xml_symbol.find("Bounds")
            x, y, width, height = bounds.get("x"), bounds.get("y"), bounds.get("w"), bounds.get("h")
            x, y, width, height = int(float(x)), int(float(y)), int(float(width)), int(float(height))

            symbol = AudiverisOmrSymbol(symbol_class, x, y, width, height)
            symbols.append(symbol)

        symbol_number = 0
        for symbol in symbols:
            symbol_class = symbol.symbol_class
            bounding_box_with_one_pixel_margin = symbol.as_bounding_box_with_margin(1)
            symbol_image = image.crop(bounding_box_with_one_pixel_margin)

            target_directory = os.path.join(destination_directory, symbol_class)
            os.makedirs(target_directory, exist_ok=True)

            export_path = ExportPath(destination_directory, symbol_class,
                                     file_name_without_extension + str(symbol_number))
            symbol_image.save(export_path.get_full_path())
            symbol_number += 1



if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--raw_dataset_directory",
        type=str,
        default="../data/audiveris_omr_raw",
        help="The directory, where the raw Audiveris OMR dataset can be found")
    parser.add_argument(
        "--image_dataset_directory",
        type=str,
        default="../data/audiveris_omr",
        help="The directory, where the generated bitmaps will be created")

    flags, unparsed = parser.parse_known_args()

    audiveris_omr_image_generator = AudiverisOmrImageGenerator()
    audiveris_omr_image_generator.extract_symbols(flags.raw_dataset_directory, flags.image_dataset_directory)