import cv2
from os import listdir
from os.path import join, exists
from os import makedirs
from lxml import etree
import png
import xml.etree.ElementTree as ET
import io
import base64

def recursivityFilling(element, file):
    for subEl in element:
        if subEl.tag == '{http://www.w3.org/2000/svg}g': # Check if leaf or node
            recursivityFilling(subEl, file)
        else:
            folder = subEl.attrib['id'].lower() # The id contains the same name as the folder
            folder = folder.replace(" ", "") # But might contain some whitespaces that we remove
            label = cv2.imread(join(INP_FOLDER, folder, file), 0)
            if a is not None:
                with io.BytesIO() as buffer:
                    pngW = png.Writer(width=label.shape[1], height=label.shape[0], greyscale=True, bitdepth=1)
                    pngW.write(buffer, label>200)
                    pngData = base64.b64encode(buffer.getvalue())
                imgAttribute = 'src="data:image/png;base64, %s"' % pngData
                subEl.set('xlink:href', imgAttribute)

list_img = listdir(join(INP_FOLDER, 'raw')) # reference of existing files.
list_img.remove('.directory')

svg_path = join(INP_FOLDER, 'svg') # creating the results folders
if not exists(svg_path):
    makedirs(svg_path)


for f in list_img:
    tree = ET.parse('fundus.svg') # The reference SVG that we use to get the expected hierarchy.
    root = tree.getroot()
    for child in root:
        recursivityFilling(child, f)
    tree.write(join(INP_FOLDER, 'svg', f[:-3]+'svg')) # Write the svg with the same name as the reference image.