"""
Skinet (Segmentation of the Kidney through a Neural nETwork) Project
Dataset tools

Copyright (c) 2021 Skinet Team
Licensed under the MIT License (see LICENSE for details)
Written by <NAME>
"""
import os
import shutil
from time import time

import cv2
import numpy as np
from skimage.io import imread

from common_utils import progressBar, formatTime
from datasetTools import AnnotationAdapter as adapt
from datasetTools.AnnotationAdapter import AnnotationAdapter
from datasetTools.datasetDivider import getBWCount, CV2_IMWRITE_PARAM
from mrcnn.Config import Config
from mrcnn.utils import extract_bboxes, expand_mask, minimize_mask, get_bboxes_intersection, shift_bbox

NEPHRO_CLASSES = [
    {"id": 0, "name": "Background", "color": "", "ignore": True},
    {"id": 1, "name": "tubule_sain", "color": "#ff007f", "ignore": False},
    {"id": 2, "name": "tubule_atrophique", "color": "#55557f", "ignore": False},
    {"id": 3, "name": "nsg_complet", "color": "#ff557f", "ignore": False},
    {"id": 4, "name": "nsg_partiel", "color": "#55aa7f", "ignore": False},
    {"id": 5, "name": "pac", "color": "#ffaa7f", "ignore": False},
    {"id": 6, "name": "vaisseau", "color": "#55ff7f", "ignore": False},
    {"id": 7, "name": "artefact", "color": "#000000", "ignore": False},
    {"id": 8, "name": "veine", "color": "#0000ff", "ignore": False},
    {"id": 9, "name": "nsg", "color": "#55007f", "ignore": False},
    {"id": 10, "name": "intima", "color": "#aa0000", "ignore": False},
    {"id": 11, "name": "media", "color": "#aa5500", "ignore": False}
]

CORTICES_CLASSES = [
    {"id": 0, "name": "Background", "color": "", "ignore": True},
    {"id": 1, "name": "cortex", "color": "#ffaa00", "ignore": False},
    {"id": 2, "name": "medullaire", "color": "#ff0000", "ignore": False},
    {"id": 3, "name": "capsule", "color": "#ff00ff", "ignore": False}
]

MESTC_GLOM_CLASSES = [
    {"id": 0, "name": "Background", "color": "", "ignore": True},
    {"id": 1, "name": "hile", "color": "#64FE2E", "ignore": False},
    {"id": 2, "name": "M", "color": "#55007f", "ignore": False},
    {"id": 3, "name": "E", "color": "#ff007f", "ignore": False},
    {"id": 4, "name": "S", "color": "#55557f", "ignore": False},
    {"id": 5, "name": "C", "color": "#ff557f", "ignore": False},
    {"id": 6, "name": "necrose_fib", "color": "#55aa7f", "ignore": False},
    {"id": 7, "name": "artefact", "color": "#ffaa7f", "ignore": False}
]

INFLAMMATION_CLASSES = [
    {"id": 0, "name": "Background", "color": "", "ignore": True},
    {"id": 1, "name": "inflammation", "color": "#ffff00", "ignore": False},
    {"id": 2, "name": "cpt", "color": "#64FE2E", "ignore": False}
]


def get_bbox_from_points(pts):
    """
    Return bbox from a points array
    :param pts: the points coordinates
    :return: y1, x1, y2, x2
    """
    temp = np.array(pts)
    x1 = np.amin(temp[:, 0])
    x2 = np.amax(temp[:, 0])
    y1 = np.amin(temp[:, 1])
    y2 = np.amax(temp[:, 1])
    return y1, x1, y2, x2


def createMask(imgName: str, imgShape, idMask: int, ptsMask, datasetName: str = 'dataset_train',
               maskClass: str = 'masks', imageFormat="jpg", config: Config = None):
    """
    Create the mask image based on its polygon points
    :param imgName: name w/o extension of the base image
    :param imgShape: shape of the image
    :param idMask: the ID of the mask, a number not already used for that image
    :param ptsMask: array of [x, y] coordinates which are all the polygon points representing the mask
    :param datasetName: name of the output dataset
    :param maskClass: name of the associated class of the current mask
    :param imageFormat: output format of the masks' images
    :param config: config object
    :return: None
    """
    # https://www.programcreek.com/python/example/89415/cv2.fillPoly
    # Formatting coordinates matrix to get int
    ptsMask = np.double(ptsMask)
    ptsMask = np.matrix.round(ptsMask)
    ptsMask = np.int32(ptsMask)

    bbox_coordinates = ""
    if config is not None and config.is_using_mini_mask():
        bbox = get_bbox_from_points(ptsMask)
        if get_bboxes_intersection(bbox, [0, 0, *imgShape[:2]]) <= 0:
            return
        kept_bbox = [0, 0, 0, 0]
        for i in range(4):
            kept_bbox[i] = min(max(0, bbox[i]), imgShape[i % 2])
        y1, x1, y2, x2 = kept_bbox
        bbox_coordinates = f"_{y1}_{x1}_{y2}_{x2}"

        shiftedBbox = shift_bbox(bbox)
        shift = bbox[:2]
        mask = np.uint8(np.zeros((shiftedBbox[2], shiftedBbox[3])))
        cv2.fillPoly(mask, [ptsMask - shift[::-1]], 255)

        shifted_kept_bbox = shift_bbox(kept_bbox, customShift=shift)
        y1, x1, y2, x2 = shifted_kept_bbox
        mask = mask[y1:y2, x1:x2]

        # Creating black matrix with same size than original image and then drawing the mask
        mask = minimize_mask(shiftedBbox, mask, config.get_mini_mask_shape())
        mask = mask.astype(np.uint8) * 255
    else:
        # Creating black matrix with same size than original image and then drawing the mask
        mask = np.uint8(np.zeros((imgShape[0], imgShape[1])))
        cv2.fillPoly(mask, [ptsMask], 255)

    # Saving result image
    maskClass = maskClass.lower().strip(' ').replace(" ", "_")
    output_directory = os.path.join(datasetName, imgName, maskClass)
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)
    output_name = f"{imgName}_{idMask:03d}{bbox_coordinates}.{imageFormat}"
    cv2.imwrite(os.path.join(output_directory, output_name), mask, CV2_IMWRITE_PARAM)


def getBboxFromName(imageName):
    """
    Return the bbox coordinates stored in the image name
    :param imageName: the image name from which you want the bbox
    :return: the bbox as [y1, x1, y2, x2]
    """
    lastParts = os.path.splitext(os.path.basename(imageName))[0].split('_')[-4:]
    return np.array([int(x) for x in lastParts])


def resizeMasks(baseMasks, xRatio: float, yRatio: float):
    """
    Resize mask's base points to fit the targeted size
    :param baseMasks array of [x, y] coordinates which are all the polygon points representing the mask
    :param xRatio width ratio that will be applied to coordinates
    :param yRatio height ratio that will be applied to coordinates
    """
    res = []
    for pt in baseMasks:
        xTemp = float(pt[0])
        yTemp = float(pt[1])
        res.append([xTemp * xRatio, yTemp * yRatio])
    return res


def createMasksOfImage(rawDatasetPath: str, imgName: str, datasetName: str = 'dataset_train',
                       adapter: AnnotationAdapter = None, classesInfo: dict = None, imageFormat="jpg", resize=None,
                       config: Config = None):
    """
    Create all the masks of a given image by parsing xml annotations file
    :param rawDatasetPath: path to the folder containing images and associated annotations
    :param imgName: name w/o extension of an image
    :param datasetName: name of the output dataset
    :param adapter: the annotation adapter to use to create masks, if None looking for an adapter that can read the file
    :param classesInfo: Information about all classes that are used, by default will be nephrology classes Info
    :param imageFormat: output format of the image and masks
    :param resize: if the image and masks have to be resized
    :param config: config object
    :return: None
    """
    # Getting shape of original image (same for all this masks)
    if classesInfo is None:
        classesInfo = NEPHRO_CLASSES if config is None else config.get_classes_info()

    img = cv2.imread(os.path.join(rawDatasetPath, f"{imgName}.{imageFormat}"))
    if img is None:
        print(f'Problem with {imgName} image')
        return
    shape = img.shape
    if resize is not None:
        yRatio = resize[0] / shape[0]
        xRatio = resize[1] / shape[1]
        assert yRatio > 0 and xRatio > 0, f"Error resize ratio not correct ({yRatio:3.2f}, {xRatio:3.2f})"
        img = cv2.resize(img, resize, interpolation=cv2.INTER_CUBIC)
        shape = img.shape

    # Copying the original image in the dataset
    targetDirectoryPath = os.path.join(datasetName, imgName, 'images')
    if not os.path.exists(targetDirectoryPath):
        os.makedirs(targetDirectoryPath)
        # TODO use file copy if unchanged else cv2
        cv2.imwrite(os.path.join(targetDirectoryPath, f"{imgName}.{imageFormat}"), img, CV2_IMWRITE_PARAM)

    # Finding annotation files
    formats = adapt.ANNOTATION_FORMAT
    fileList = os.listdir(rawDatasetPath)
    imageFiles = []
    for file in fileList:
        if imgName in file:
            if file.split('.')[-1] in formats:
                imageFiles.append(file)

    # Choosing the adapter to use (parameters to force it ?)
    file = None
    assert len(imageFiles) > 0
    if adapter is None:
        # No adapter given, we are looking for the adapter with highest priority level that can read an/the annotation
        # file
        adapters = list(adapt.ANNOTATION_ADAPTERS.values())
        adapterPriority = -1
        for f in imageFiles:
            for a in adapters:
                if a.canRead(os.path.join(rawDatasetPath, f)):
                    if a.getPriorityLevel() > adapterPriority:
                        adapterPriority = a.getPriorityLevel()
                        adapter = a
                        file = f
    else:
        # Using given adapter, we are looking for a file that can be read
        file = None
        for f in imageFiles:
            if adapter.canRead(os.path.join(rawDatasetPath, f)) and file is None:
                file = f

    # Getting the masks data
    masks = adapter.readFile(os.path.join(rawDatasetPath, file))

    # Creating masks
    for noMask, (datasetClass, maskPoints) in enumerate(masks):
        # Converting class id to class name if needed
        if type(datasetClass) is int:
            if datasetClass < len(classesInfo) and classesInfo[datasetClass]["id"] == datasetClass:
                maskClass = classesInfo[datasetClass]["name"]
            else:
                for classInfo in classesInfo:
                    if classInfo["id"] == datasetClass:
                        maskClass = classInfo["name"]
                        break
        else:
            maskClass = datasetClass
            if maskClass == "None":
                print(f" /!\\ {imgName} : None class present /!\\ ")
        if resize is not None:
            resizedMasks = resizeMasks(maskPoints, xRatio, yRatio)
        createMask(imgName, shape, noMask, maskPoints if resize is None else resizedMasks, datasetName, maskClass,
                   imageFormat, config=config)


def fuseClassMasks(datasetPath: str, imageName: str, targetedClass: [str, list], imageFormat="jpg",
                   deleteBaseMasks=False, silent=False):
    """
    Fuse each targeted-class masks into one
    :param datasetPath: the dataset that have been wrapped
    :param imageName: the image you want its targeted-class to be fused
    :param targetedClass: the class of the masks that have to be fused
    :param imageFormat: format to use to save the final targeted-class masks
    :param deleteBaseMasks: delete the base masks images after fusion
    :param silent: if True will not print
    :return: None
    """
    if type(targetedClass) is str:
        targetedClasses = [targetedClass]
    else:
        targetedClasses = targetedClass
    # Getting the image directory path
    imageDir = os.path.join(datasetPath, imageName)
    imagePath = os.path.join(datasetPath, imageName, "images")
    imagePath = os.path.join(imagePath, os.listdir(imagePath)[0])
    image = cv2.imread(imagePath)
    for aClass in targetedClasses:
        classDir = os.path.join(imageDir, aClass)
        if os.path.exists(classDir):
            listClassImages = os.listdir(classDir)
            if not silent:
                print(f"Fusing {imageName} {aClass} class masks")
            fusion = loadSameResImage(os.path.join(classDir, listClassImages[0]), imageShape=image.shape)
            listClassImages.remove(listClassImages[0])
            for maskName in listClassImages:  # Adding each mask to the same image
                maskPath = os.path.join(classDir, maskName)
                mask = loadSameResImage(maskPath, imageShape=image.shape)
                fusion = cv2.add(fusion, mask)
            # Saving the fused mask image
            cv2.imwrite(os.path.join(classDir, f"{imageName}_{aClass}.{imageFormat}"), fusion, CV2_IMWRITE_PARAM)
            if deleteBaseMasks:
                for maskName in os.listdir(classDir):  # Deleting each targeted-class mask except the fused one
                    if f'_{aClass}.{imageFormat}' not in maskName:
                        maskPath = os.path.join(classDir, maskName)
                        os.remove(maskPath)


def cleanFusedClassDir(datasetPath: str, fusedClass):
    """
    Cleaning all fused-class directories in the dataset, keeping only unique file or fused ones
    :param datasetPath: the dataset that have been wrapped
    :param fusedClass: the class 
    :return: None
    """
    for imageDir in os.listdir(datasetPath):
        imageDirPath = os.path.join(datasetPath, imageDir)
        fusedClassDirPath = os.path.join(imageDirPath, fusedClass)
        if os.path.exists(fusedClassDirPath):
            listFusedClassImages = os.listdir(fusedClassDirPath)
            if len(listFusedClassImages) > 1:  # Deleting only if strictly more than one fused-class mask image present
                fusedFusedClassPresent = False
                for fusedClassImage in listFusedClassImages:
                    fusedFusedClassPresent = fusedFusedClassPresent or (f'_{fusedClass}' in fusedClassImage)
                if fusedFusedClassPresent:
                    # Deleting each fused-class mask except the fused one
                    for maskName in os.listdir(fusedClassDirPath):
                        if f'_{fusedClass}' not in maskName:
                            maskPath = os.path.join(fusedClassDirPath, maskName)
                            os.remove(maskPath)


def cleanImage(datasetPath: str, imageName: str, cleaningClasses: str, excludeClasses=None, imageFormat="jpg",
               cleanMasks=False, minAreaThreshold=300, config: Config = None):
    """
    Creating the full_images directory and cleaning the base image by removing non-cleaning-class areas
    :param excludeClasses:
    :param datasetPath: the dataset that have been wrapped
    :param imageName: the image name
    :param cleaningClasses: the class to use to clean the image
    :param cleanMasks: if true, will clean masks based on the cleaning-class-mask
    :param imageFormat: the image format to use to save the image
    :param minAreaThreshold: remove mask if its area is smaller than this threshold
    :param config: config object
    :return: None
    """
    assert cleaningClasses is not None and cleaningClasses != "", "Cleaning class is required."
    if type(cleaningClasses) is str:
        cleaningClasses = [cleaningClasses]
    if type(excludeClasses) is str:
        excludeClasses = [excludeClasses]

    # Getting the base image
    path = os.path.join(datasetPath, imageName, '{folder}', f"{imageName}.{imageFormat}")
    imagePath = path.format(folder='images')
    fullImagePath = path.format(folder='full_images')
    image = cv2.imread(imagePath)

    # Fusing all the cleaning-class masks and then cleaning the image and if needed the masks
    cleaningClassMasks = gatherClassesMasks(datasetPath, imageName, image.shape, cleaningClasses)
    if excludeClasses is None:
        excludedClassMasks = None
    else:
        excludedClassMasks = gatherClassesMasks(datasetPath, imageName, image.shape, excludeClasses)

    if cleaningClassMasks is not None or excludedClassMasks is not None:
        if cleaningClassMasks is None:
            cleaningClassMasks = np.ones_like(image)[..., 0] * 255
        if excludedClassMasks is not None:
            excludedClassMasks = cv2.bitwise_not(excludedClassMasks)
            cleaningClassMasks = cv2.bitwise_and(cleaningClassMasks, excludedClassMasks)
        # Copying the full image into the correct directory
        os.makedirs(os.path.dirname(fullImagePath), exist_ok=True)
        shutil.copy2(imagePath, fullImagePath)

        # Cleaning the image and saving it
        image = cv2.bitwise_and(image, np.repeat(cleaningClassMasks[:, :, np.newaxis], 3, axis=2))
        cv2.imwrite(imagePath, image, CV2_IMWRITE_PARAM)

        # Cleaning masks so that they cannot exist elsewhere
        if cleanMasks:
            folderToRemove = []
            for folder in os.listdir(os.path.join(datasetPath, imageName)):
                folderPath = os.path.join(datasetPath, imageName, folder)
                # Checking only for the other classes folder
                skipClasses = ["images", "full_images"]
                skipClasses.extend(cleaningClasses)
                skipClasses.extend(excludeClasses)
                if os.path.isdir(folderPath) and folder not in skipClasses:
                    # For each mask of the folder
                    for maskImageFileName in os.listdir(folderPath):
                        maskImagePath = os.path.join(folderPath, maskImageFileName)
                        mask = loadSameResImage(maskImagePath, image.shape)
                        areaBefore = getBWCount(mask)[1]

                        # If mask is not empty
                        if areaBefore > 0:
                            # Cleaning it with the cleaning-class masks
                            mask = cv2.bitwise_and(mask, cleaningClassMasks)
                            areaAfter = getBWCount(mask)[1]
                        else:
                            areaAfter = areaBefore

                        # If mask was empty or too small after cleaning, we remove it
                        if areaBefore == 0 or areaAfter < minAreaThreshold:
                            os.remove(maskImagePath)
                        elif areaBefore != areaAfter:
                            # If mask has is different after cleaning, we replace the original one
                            try:
                                try:
                                    idMask = int(maskImageFileName.split('.')[0].split('_')[1])
                                except ValueError:
                                    # If we could not retrieve the original mask ID, give it a unique one
                                    idMask = int(time())

                                # If mini-mask are enabled, we minimize it before saving it
                                bbox_coordinates = ""
                                if config is not None and config.is_using_mini_mask():
                                    bbox = extract_bboxes(mask)
                                    mask = minimize_mask(bbox, mask, config.get_mini_mask_shape())
                                    mask = mask.astype(np.uint8) * 255
                                    y1, x1, y2, x2 = bbox
                                    bbox_coordinates = f"_{y1}_{x1}_{y2}_{x2}"

                                # Saving cleaned mask
                                outputName = f"{imageName}_{idMask:03d}{bbox_coordinates}.{imageFormat}"
                                cv2.imwrite(os.path.join(folderPath, outputName), mask, CV2_IMWRITE_PARAM)
                                if outputName != maskImageFileName:  # Remove former mask if not the same name
                                    os.remove(maskImagePath)
                            except Exception:
                                print(f"Error on {maskImagePath} update")

                    if len(os.listdir(folderPath)) == 0:
                        folderToRemove.append(folderPath)
            for folderPath in folderToRemove:
                shutil.rmtree(folderPath, ignore_errors=True)
            pass


def gatherClassesMasks(datasetPath, imageName, img_shape, gatheredClasses):
    if type(gatheredClasses) is str:
        gatheredClasses = [gatheredClasses]
    gatheredMaskFound = False
    gatheredClassMasks = None
    for gatheredClass_ in gatheredClasses:
        gatheredClassDirPath = os.path.join(datasetPath, imageName, gatheredClass_)
        gatheredClassExists = os.path.exists(gatheredClassDirPath)
        if gatheredClassExists and len(os.listdir(gatheredClassDirPath)) > 0:
            for gatheredClassMaskName in os.listdir(gatheredClassDirPath):
                gatheredClassMaskPath = os.path.join(gatheredClassDirPath, gatheredClassMaskName)
                if not gatheredMaskFound:
                    gatheredClassMasks = loadSameResImage(gatheredClassMaskPath, img_shape)
                    gatheredMaskFound = True
                else:  # Adding additional masks
                    temp, tempBbox = loadOnlyMask(gatheredClassMaskPath, img_shape)
                    y1, x1, y2, x2 = tempBbox
                    gatheredClassMasks[y1:y2, x1:x2] = cv2.bitwise_or(gatheredClassMasks[y1:y2, x1:x2], temp)
                    del temp
    return gatheredClassMasks


def loadSameResImage(imagePath, imageShape):
    mask = cv2.imread(imagePath, cv2.IMREAD_UNCHANGED)
    if mask.shape[0] != imageShape[0] or mask.shape[1] != imageShape[1]:
        bbox = getBboxFromName(imagePath)
        mask = expand_mask(bbox, mask, image_shape=imageShape)
        mask = mask.astype(np.uint8) * 255
    return mask


def loadOnlyMask(imagePath, imageShape):
    mask = cv2.imread(imagePath, cv2.IMREAD_UNCHANGED)
    if mask.shape[0] != imageShape[0] or mask.shape[1] != imageShape[1]:
        # Finding bbox coordinates from image name
        bbox = getBboxFromName(imagePath)
        shifted = shift_bbox(bbox)
        y1, x1, y2, x2 = shifted

        # Expanding mask to its real size
        mask = expand_mask(shifted, mask, image_shape=shifted[2:])
        mask = mask.astype(np.uint8) * 255
    else:
        # Extracting bbox of
        bbox = extract_bboxes(mask)
        y1, x1, y2, x2 = bbox
    return mask[y1:y2, x1:x2, ...], bbox


def convertImage(inputImagePath: str, outputImagePath: str):
    """
    Convert an image from a format to another one
    :param inputImagePath: path to the initial image
    :param outputImagePath: path to the output image
    :return: None
    """
    image = imread(inputImagePath)
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    cv2.imwrite(outputImagePath, image, CV2_IMWRITE_PARAM)


def getInfoRawDataset(rawDatasetPath: str, verbose=False, adapter: AnnotationAdapter = None, mainFormat="jpg"):
    """
    Listing all available images, those with missing information
    :param verbose: whether or not print should be executed
    :param rawDatasetPath: path to the raw dataset folder
    :param adapter:
    :param mainFormat: the format to use for the dataset
    :return: list of unique files names, list of available images names, list of missing images names,
    list of missing annotations names
    """
    names = []
    images = []  # list of image that can be used to compute masks
    missingImages = []  # list of missing images
    missingAnnotations = []  # list of missing annotations
    inputFormats = ["jpg", "jp2", "png"]
    if adapter is None:
        annotationFormats = adapt.ANNOTATION_FORMAT
    else:
        annotationFormats = [adapter.getAnnotationFormat()]
    fileList = os.listdir(rawDatasetPath)
    if verbose:
        progressBar(0, len(fileList), "Listing files")
    for idx, file in enumerate(fileList):
        if verbose and (idx % 10 == 0 or idx + 1 == len(fileList)):
            lastIdx = idx
            progressBar(idx + 1, len(fileList), "Listing files")
        name = file.split('.')[0]
        if name not in names:  # We want to do this only once per unique file name (without extension)
            names.append(name)

            availableFormat = []
            for format in inputFormats:
                imgPath = os.path.join(rawDatasetPath, f"{name}.{format}")
                if os.path.exists(imgPath):
                    availableFormat.append(format)

            # Testing if annotation file exists for that name
            annotationsExist = False
            for ext in annotationFormats:
                annotationsExist = annotationsExist or os.path.exists(os.path.join(rawDatasetPath, f"{name}.{ext}"))
            if len(availableFormat) > 0:  # At least one image exists
                if not annotationsExist:  # Annotations are missing
                    missingAnnotations.append(name)
                else:
                    if mainFormat not in availableFormat:
                        for format in inputFormats:
                            if format in availableFormat:
                                sourcePath = os.path.join(rawDatasetPath, f"{name}.{format}")
                                outputPath = os.path.join(rawDatasetPath, f"{name}.{mainFormat}")
                                convertImage(sourcePath, outputPath)
                                break
                    images.append(name)  # Adding this image to the list
            elif annotationsExist:  # There is no image file but xml found
                missingImages.append(name)
    if verbose:
        if lastIdx + 1 != len(fileList):
            progressBar(1, 1, "Listing files")
        # Displaying missing image files
        problem = False
        nbMissingImg = len(missingImages)
        if nbMissingImg > 0:
            problem = True
            print('Missing {} image{} : {}'.format(nbMissingImg, 's' if nbMissingImg > 1 else '', missingImages))

        # Displaying missing annotations files
        nbMissingAnnotations = len(missingAnnotations)
        if nbMissingAnnotations > 0:
            problem = True
            print('Missing {} annotation{} : {}'.format(nbMissingAnnotations, 's' if nbMissingAnnotations > 1 else '',
                                                        missingAnnotations))

        # Checking if there is file that is not image nor annotation
        nbImages = len(images)
        if len(names) - nbMissingImg - nbMissingAnnotations - nbImages != 0:
            problem = True
            print('Be careful, there are not only required dataset files in this folder')

        if not problem:
            print("Raw Dataset has no problem. Number of Images : {}".format(nbImages))
    return names, images, missingImages, missingAnnotations


def startWrapper(rawDatasetPath: str, datasetName: str = 'dataset_train', deleteBaseMasks=False,
                 adapter: AnnotationAdapter = None, resize=None, mode="main", classesInfo=None, imageFormat="jpg"):
    """
    Start wrapping the raw dataset into the wanted format
    :param rawDatasetPath: path to the folder containing images and associated annotations
    :param datasetName: name of the output dataset
    :param deleteBaseMasks: delete the base masks images after fusion
    :param adapter: Adapter to use to read annotations, if None compatible adapter will be searched
    :param resize: If tuple given, the images and their masks will be resized to the tuple value
    :param mode: Mode to use
    :param classesInfo: Information about the classes that will be used
    :param imageFormat: the image format to use in the dataset
    :return: None
    """
    names, images, missingImages, missingAnnotations = getInfoRawDataset(rawDatasetPath, verbose=True,
                                                                         adapter=adapter, mainFormat=imageFormat)
    if classesInfo is None:
        if mode == 'main':
            classesInfo = NEPHRO_CLASSES
        elif mode == "cortex":
            classesInfo = CORTICES_CLASSES
        elif mode == "mest_glom":
            classesInfo = MESTC_GLOM_CLASSES
        elif mode == "inflammation":
            classesInfo = INFLAMMATION_CLASSES
    # Creating masks for any image which has all required files and displaying progress
    start_time = time()
    for index, file in enumerate(images):
        progressBar(index, len(images), prefix='Creating masks',
                    suffix=f" {formatTime(round(time() - start_time))} Current : {file}")
        createMasksOfImage(rawDatasetPath, file, datasetName, adapter, classesInfo=classesInfo,
                           resize=resize, imageFormat=imageFormat)
        if mode == "main":
            fuseClassMasks(datasetName, file, "cortex", deleteBaseMasks=deleteBaseMasks, silent=True)
            cleanImage(datasetName, file, cleaningClasses='cortex')
        elif mode == "mest_glom":
            cleanImage(datasetName, file, cleaningClasses='nsg', cleanMasks=True)
        elif mode == "inflammation":
            cleanImage(datasetName, file, cleaningClasses='cortex', cleanMasks=True,
                       excludeClasses=[c["name"] for c in NEPHRO_CLASSES if not c["ignore"]])
    progressBar(1, 1, prefix='Creating masks', suffix=f"{formatTime(round(time() - start_time))}" + " " * 25)

"""
Skinet (Segmentation of the Kidney through a Neural nETwork) Project
Dataset tools

Copyright (c) 2021 Skinet Team
Licensed under the MIT License (see LICENSE for details)
Written by <NAME>
"""
import os
import shutil
from time import time

import cv2
import numpy as np
from skimage.io import imread

from common_utils import progressBar, formatTime
from datasetTools import AnnotationAdapter as adapt
from datasetTools.AnnotationAdapter import AnnotationAdapter
from datasetTools.datasetDivider import getBWCount, CV2_IMWRITE_PARAM
from mrcnn.Config import Config
from mrcnn.utils import extract_bboxes, expand_mask, minimize_mask, get_bboxes_intersection, shift_bbox

NEPHRO_CLASSES = [
    {"id": 0, "name": "Background", "color": "", "ignore": True},
    {"id": 1, "name": "tubule_sain", "color": "#ff007f", "ignore": False},
    {"id": 2, "name": "tubule_atrophique", "color": "#55557f", "ignore": False},
    {"id": 3, "name": "nsg_complet", "color": "#ff557f", "ignore": False},
    {"id": 4, "name": "nsg_partiel", "color": "#55aa7f", "ignore": False},
    {"id": 5, "name": "pac", "color": "#ffaa7f", "ignore": False},
    {"id": 6, "name": "vaisseau", "color": "#55ff7f", "ignore": False},
    {"id": 7, "name": "artefact", "color": "#000000", "ignore": False},
    {"id": 8, "name": "veine", "color": "#0000ff", "ignore": False},
    {"id": 9, "name": "nsg", "color": "#55007f", "ignore": False},
    {"id": 10, "name": "intima", "color": "#aa0000", "ignore": False},
    {"id": 11, "name": "media", "color": "#aa5500", "ignore": False}
]

CORTICES_CLASSES = [
    {"id": 0, "name": "Background", "color": "", "ignore": True},
    {"id": 1, "name": "cortex", "color": "#ffaa00", "ignore": False},
    {"id": 2, "name": "medullaire", "color": "#ff0000", "ignore": False},
    {"id": 3, "name": "capsule", "color": "#ff00ff", "ignore": False}
]

MESTC_GLOM_CLASSES = [
    {"id": 0, "name": "Background", "color": "", "ignore": True},
    {"id": 1, "name": "hile", "color": "#64FE2E", "ignore": False},
    {"id": 2, "name": "M", "color": "#55007f", "ignore": False},
    {"id": 3, "name": "E", "color": "#ff007f", "ignore": False},
    {"id": 4, "name": "S", "color": "#55557f", "ignore": False},
    {"id": 5, "name": "C", "color": "#ff557f", "ignore": False},
    {"id": 6, "name": "necrose_fib", "color": "#55aa7f", "ignore": False},
    {"id": 7, "name": "artefact", "color": "#ffaa7f", "ignore": False}
]

INFLAMMATION_CLASSES = [
    {"id": 0, "name": "Background", "color": "", "ignore": True},
    {"id": 1, "name": "inflammation", "color": "#ffff00", "ignore": False},
    {"id": 2, "name": "cpt", "color": "#64FE2E", "ignore": False}
]


def get_bbox_from_points(pts):
    """
    Return bbox from a points array
    :param pts: the points coordinates
    :return: y1, x1, y2, x2
    """
    temp = np.array(pts)
    x1 = np.amin(temp[:, 0])
    x2 = np.amax(temp[:, 0])
    y1 = np.amin(temp[:, 1])
    y2 = np.amax(temp[:, 1])
    return y1, x1, y2, x2


def createMask(imgName: str, imgShape, idMask: int, ptsMask, datasetName: str = 'dataset_train',
               maskClass: str = 'masks', imageFormat="jpg", config: Config = None):
    """
    Create the mask image based on its polygon points
    :param imgName: name w/o extension of the base image
    :param imgShape: shape of the image
    :param idMask: the ID of the mask, a number not already used for that image
    :param ptsMask: array of [x, y] coordinates which are all the polygon points representing the mask
    :param datasetName: name of the output dataset
    :param maskClass: name of the associated class of the current mask
    :param imageFormat: output format of the masks' images
    :param config: config object
    :return: None
    """
    # https://www.programcreek.com/python/example/89415/cv2.fillPoly
    # Formatting coordinates matrix to get int
    ptsMask = np.double(ptsMask)
    ptsMask = np.matrix.round(ptsMask)
    ptsMask = np.int32(ptsMask)

    bbox_coordinates = ""
    if config is not None and config.is_using_mini_mask():
        bbox = get_bbox_from_points(ptsMask)
        if get_bboxes_intersection(bbox, [0, 0, *imgShape[:2]]) <= 0:
            return
        kept_bbox = [0, 0, 0, 0]
        for i in range(4):
            kept_bbox[i] = min(max(0, bbox[i]), imgShape[i % 2])
        y1, x1, y2, x2 = kept_bbox
        bbox_coordinates = f"_{y1}_{x1}_{y2}_{x2}"

        shiftedBbox = shift_bbox(bbox)
        shift = bbox[:2]
        mask = np.uint8(np.zeros((shiftedBbox[2], shiftedBbox[3])))
        cv2.fillPoly(mask, [ptsMask - shift[::-1]], 255)

        shifted_kept_bbox = shift_bbox(kept_bbox, customShift=shift)
        y1, x1, y2, x2 = shifted_kept_bbox
        mask = mask[y1:y2, x1:x2]

        # Creating black matrix with same size than original image and then drawing the mask
        mask = minimize_mask(shiftedBbox, mask, config.get_mini_mask_shape())
        mask = mask.astype(np.uint8) * 255
    else:
        # Creating black matrix with same size than original image and then drawing the mask
        mask = np.uint8(np.zeros((imgShape[0], imgShape[1])))
        cv2.fillPoly(mask, [ptsMask], 255)

    # Saving result image
    maskClass = maskClass.lower().strip(' ').replace(" ", "_")
    output_directory = os.path.join(datasetName, imgName, maskClass)
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)
    output_name = f"{imgName}_{idMask:03d}{bbox_coordinates}.{imageFormat}"
    cv2.imwrite(os.path.join(output_directory, output_name), mask, CV2_IMWRITE_PARAM)


def getBboxFromName(imageName):
    """
    Return the bbox coordinates stored in the image name
    :param imageName: the image name from which you want the bbox
    :return: the bbox as [y1, x1, y2, x2]
    """
    lastParts = os.path.splitext(os.path.basename(imageName))[0].split('_')[-4:]
    return np.array([int(x) for x in lastParts])


def resizeMasks(baseMasks, xRatio: float, yRatio: float):
    """
    Resize mask's base points to fit the targeted size
    :param baseMasks array of [x, y] coordinates which are all the polygon points representing the mask
    :param xRatio width ratio that will be applied to coordinates
    :param yRatio height ratio that will be applied to coordinates
    """
    res = []
    for pt in baseMasks:
        xTemp = float(pt[0])
        yTemp = float(pt[1])
        res.append([xTemp * xRatio, yTemp * yRatio])
    return res


def createMasksOfImage(rawDatasetPath: str, imgName: str, datasetName: str = 'dataset_train',
                       adapter: AnnotationAdapter = None, classesInfo: dict = None, imageFormat="jpg", resize=None,
                       config: Config = None):
    """
    Create all the masks of a given image by parsing xml annotations file
    :param rawDatasetPath: path to the folder containing images and associated annotations
    :param imgName: name w/o extension of an image
    :param datasetName: name of the output dataset
    :param adapter: the annotation adapter to use to create masks, if None looking for an adapter that can read the file
    :param classesInfo: Information about all classes that are used, by default will be nephrology classes Info
    :param imageFormat: output format of the image and masks
    :param resize: if the image and masks have to be resized
    :param config: config object
    :return: None
    """
    # Getting shape of original image (same for all this masks)
    if classesInfo is None:
        classesInfo = NEPHRO_CLASSES if config is None else config.get_classes_info()

    img = cv2.imread(os.path.join(rawDatasetPath, f"{imgName}.{imageFormat}"))
    if img is None:
        print(f'Problem with {imgName} image')
        return
    shape = img.shape
    if resize is not None:
        yRatio = resize[0] / shape[0]
        xRatio = resize[1] / shape[1]
        assert yRatio > 0 and xRatio > 0, f"Error resize ratio not correct ({yRatio:3.2f}, {xRatio:3.2f})"
        img = cv2.resize(img, resize, interpolation=cv2.INTER_CUBIC)
        shape = img.shape

    # Copying the original image in the dataset
    targetDirectoryPath = os.path.join(datasetName, imgName, 'images')
    if not os.path.exists(targetDirectoryPath):
        os.makedirs(targetDirectoryPath)
        # TODO use file copy if unchanged else cv2
        cv2.imwrite(os.path.join(targetDirectoryPath, f"{imgName}.{imageFormat}"), img, CV2_IMWRITE_PARAM)

    # Finding annotation files
    formats = adapt.ANNOTATION_FORMAT
    fileList = os.listdir(rawDatasetPath)
    imageFiles = []
    for file in fileList:
        if imgName in file:
            if file.split('.')[-1] in formats:
                imageFiles.append(file)

    # Choosing the adapter to use (parameters to force it ?)
    file = None
    assert len(imageFiles) > 0
    if adapter is None:
        # No adapter given, we are looking for the adapter with highest priority level that can read an/the annotation
        # file
        adapters = list(adapt.ANNOTATION_ADAPTERS.values())
        adapterPriority = -1
        for f in imageFiles:
            for a in adapters:
                if a.canRead(os.path.join(rawDatasetPath, f)):
                    if a.getPriorityLevel() > adapterPriority:
                        adapterPriority = a.getPriorityLevel()
                        adapter = a
                        file = f
    else:
        # Using given adapter, we are looking for a file that can be read
        file = None
        for f in imageFiles:
            if adapter.canRead(os.path.join(rawDatasetPath, f)) and file is None:
                file = f

    # Getting the masks data
    masks = adapter.readFile(os.path.join(rawDatasetPath, file))

    # Creating masks
    for noMask, (datasetClass, maskPoints) in enumerate(masks):
        # Converting class id to class name if needed
        if type(datasetClass) is int:
            if datasetClass < len(classesInfo) and classesInfo[datasetClass]["id"] == datasetClass:
                maskClass = classesInfo[datasetClass]["name"]
            else:
                for classInfo in classesInfo:
                    if classInfo["id"] == datasetClass:
                        maskClass = classInfo["name"]
                        break
        else:
            maskClass = datasetClass
            if maskClass == "None":
                print(f" /!\\ {imgName} : None class present /!\\ ")
        if resize is not None:
            resizedMasks = resizeMasks(maskPoints, xRatio, yRatio)
        createMask(imgName, shape, noMask, maskPoints if resize is None else resizedMasks, datasetName, maskClass,
                   imageFormat, config=config)


def fuseClassMasks(datasetPath: str, imageName: str, targetedClass: [str, list], imageFormat="jpg",
                   deleteBaseMasks=False, silent=False):
    """
    Fuse each targeted-class masks into one
    :param datasetPath: the dataset that have been wrapped
    :param imageName: the image you want its targeted-class to be fused
    :param targetedClass: the class of the masks that have to be fused
    :param imageFormat: format to use to save the final targeted-class masks
    :param deleteBaseMasks: delete the base masks images after fusion
    :param silent: if True will not print
    :return: None
    """
    if type(targetedClass) is str:
        targetedClasses = [targetedClass]
    else:
        targetedClasses = targetedClass
    # Getting the image directory path
    imageDir = os.path.join(datasetPath, imageName)
    imagePath = os.path.join(datasetPath, imageName, "images")
    imagePath = os.path.join(imagePath, os.listdir(imagePath)[0])
    image = cv2.imread(imagePath)
    for aClass in targetedClasses:
        classDir = os.path.join(imageDir, aClass)
        if os.path.exists(classDir):
            listClassImages = os.listdir(classDir)
            if not silent:
                print(f"Fusing {imageName} {aClass} class masks")
            fusion = loadSameResImage(os.path.join(classDir, listClassImages[0]), imageShape=image.shape)
            listClassImages.remove(listClassImages[0])
            for maskName in listClassImages:  # Adding each mask to the same image
                maskPath = os.path.join(classDir, maskName)
                mask = loadSameResImage(maskPath, imageShape=image.shape)
                fusion = cv2.add(fusion, mask)
            # Saving the fused mask image
            cv2.imwrite(os.path.join(classDir, f"{imageName}_{aClass}.{imageFormat}"), fusion, CV2_IMWRITE_PARAM)
            if deleteBaseMasks:
                for maskName in os.listdir(classDir):  # Deleting each targeted-class mask except the fused one
                    if f'_{aClass}.{imageFormat}' not in maskName:
                        maskPath = os.path.join(classDir, maskName)
                        os.remove(maskPath)


def cleanFusedClassDir(datasetPath: str, fusedClass):
    """
    Cleaning all fused-class directories in the dataset, keeping only unique file or fused ones
    :param datasetPath: the dataset that have been wrapped
    :param fusedClass: the class 
    :return: None
    """
    for imageDir in os.listdir(datasetPath):
        imageDirPath = os.path.join(datasetPath, imageDir)
        fusedClassDirPath = os.path.join(imageDirPath, fusedClass)
        if os.path.exists(fusedClassDirPath):
            listFusedClassImages = os.listdir(fusedClassDirPath)
            if len(listFusedClassImages) > 1:  # Deleting only if strictly more than one fused-class mask image present
                fusedFusedClassPresent = False
                for fusedClassImage in listFusedClassImages:
                    fusedFusedClassPresent = fusedFusedClassPresent or (f'_{fusedClass}' in fusedClassImage)
                if fusedFusedClassPresent:
                    # Deleting each fused-class mask except the fused one
                    for maskName in os.listdir(fusedClassDirPath):
                        if f'_{fusedClass}' not in maskName:
                            maskPath = os.path.join(fusedClassDirPath, maskName)
                            os.remove(maskPath)


def cleanImage(datasetPath: str, imageName: str, cleaningClasses: str, excludeClasses=None, imageFormat="jpg",
               cleanMasks=False, minAreaThreshold=300, config: Config = None):
    """
    Creating the full_images directory and cleaning the base image by removing non-cleaning-class areas
    :param excludeClasses:
    :param datasetPath: the dataset that have been wrapped
    :param imageName: the image name
    :param cleaningClasses: the class to use to clean the image
    :param cleanMasks: if true, will clean masks based on the cleaning-class-mask
    :param imageFormat: the image format to use to save the image
    :param minAreaThreshold: remove mask if its area is smaller than this threshold
    :param config: config object
    :return: None
    """
    assert cleaningClasses is not None and cleaningClasses != "", "Cleaning class is required."
    if type(cleaningClasses) is str:
        cleaningClasses = [cleaningClasses]
    if type(excludeClasses) is str:
        excludeClasses = [excludeClasses]

    # Getting the base image
    path = os.path.join(datasetPath, imageName, '{folder}', f"{imageName}.{imageFormat}")
    imagePath = path.format(folder='images')
    fullImagePath = path.format(folder='full_images')
    image = cv2.imread(imagePath)

    # Fusing all the cleaning-class masks and then cleaning the image and if needed the masks
    cleaningClassMasks = gatherClassesMasks(datasetPath, imageName, image.shape, cleaningClasses)
    if excludeClasses is None:
        excludedClassMasks = None
    else:
        excludedClassMasks = gatherClassesMasks(datasetPath, imageName, image.shape, excludeClasses)

    if cleaningClassMasks is not None or excludedClassMasks is not None:
        if cleaningClassMasks is None:
            cleaningClassMasks = np.ones_like(image)[..., 0] * 255
        if excludedClassMasks is not None:
            excludedClassMasks = cv2.bitwise_not(excludedClassMasks)
            cleaningClassMasks = cv2.bitwise_and(cleaningClassMasks, excludedClassMasks)
        # Copying the full image into the correct directory
        os.makedirs(os.path.dirname(fullImagePath), exist_ok=True)
        shutil.copy2(imagePath, fullImagePath)

        # Cleaning the image and saving it
        image = cv2.bitwise_and(image, np.repeat(cleaningClassMasks[:, :, np.newaxis], 3, axis=2))
        cv2.imwrite(imagePath, image, CV2_IMWRITE_PARAM)

        # Cleaning masks so that they cannot exist elsewhere
        if cleanMasks:
            folderToRemove = []
            for folder in os.listdir(os.path.join(datasetPath, imageName)):
                folderPath = os.path.join(datasetPath, imageName, folder)
                # Checking only for the other classes folder
                skipClasses = ["images", "full_images"]
                skipClasses.extend(cleaningClasses)
                skipClasses.extend(excludeClasses)
                if os.path.isdir(folderPath) and folder not in skipClasses:
                    # For each mask of the folder
                    for maskImageFileName in os.listdir(folderPath):
                        maskImagePath = os.path.join(folderPath, maskImageFileName)
                        mask = loadSameResImage(maskImagePath, image.shape)
                        areaBefore = getBWCount(mask)[1]

                        # If mask is not empty
                        if areaBefore > 0:
                            # Cleaning it with the cleaning-class masks
                            mask = cv2.bitwise_and(mask, cleaningClassMasks)
                            areaAfter = getBWCount(mask)[1]
                        else:
                            areaAfter = areaBefore

                        # If mask was empty or too small after cleaning, we remove it
                        if areaBefore == 0 or areaAfter < minAreaThreshold:
                            os.remove(maskImagePath)
                        elif areaBefore != areaAfter:
                            # If mask has is different after cleaning, we replace the original one
                            try:
                                try:
                                    idMask = int(maskImageFileName.split('.')[0].split('_')[1])
                                except ValueError:
                                    # If we could not retrieve the original mask ID, give it a unique one
                                    idMask = int(time())

                                # If mini-mask are enabled, we minimize it before saving it
                                bbox_coordinates = ""
                                if config is not None and config.is_using_mini_mask():
                                    bbox = extract_bboxes(mask)
                                    mask = minimize_mask(bbox, mask, config.get_mini_mask_shape())
                                    mask = mask.astype(np.uint8) * 255
                                    y1, x1, y2, x2 = bbox
                                    bbox_coordinates = f"_{y1}_{x1}_{y2}_{x2}"

                                # Saving cleaned mask
                                outputName = f"{imageName}_{idMask:03d}{bbox_coordinates}.{imageFormat}"
                                cv2.imwrite(os.path.join(folderPath, outputName), mask, CV2_IMWRITE_PARAM)
                                if outputName != maskImageFileName:  # Remove former mask if not the same name
                                    os.remove(maskImagePath)
                            except Exception:
                                print(f"Error on {maskImagePath} update")

                    if len(os.listdir(folderPath)) == 0:
                        folderToRemove.append(folderPath)
            for folderPath in folderToRemove:
                shutil.rmtree(folderPath, ignore_errors=True)
            pass


def gatherClassesMasks(datasetPath, imageName, img_shape, gatheredClasses):
    if type(gatheredClasses) is str:
        gatheredClasses = [gatheredClasses]
    gatheredMaskFound = False
    gatheredClassMasks = None
    for gatheredClass_ in gatheredClasses:
        gatheredClassDirPath = os.path.join(datasetPath, imageName, gatheredClass_)
        gatheredClassExists = os.path.exists(gatheredClassDirPath)
        if gatheredClassExists and len(os.listdir(gatheredClassDirPath)) > 0:
            for gatheredClassMaskName in os.listdir(gatheredClassDirPath):
                gatheredClassMaskPath = os.path.join(gatheredClassDirPath, gatheredClassMaskName)
                if not gatheredMaskFound:
                    gatheredClassMasks = loadSameResImage(gatheredClassMaskPath, img_shape)
                    gatheredMaskFound = True
                else:  # Adding additional masks
                    temp, tempBbox = loadOnlyMask(gatheredClassMaskPath, img_shape)
                    y1, x1, y2, x2 = tempBbox
                    gatheredClassMasks[y1:y2, x1:x2] = cv2.bitwise_or(gatheredClassMasks[y1:y2, x1:x2], temp)
                    del temp
    return gatheredClassMasks


def loadSameResImage(imagePath, imageShape):
    mask = cv2.imread(imagePath, cv2.IMREAD_UNCHANGED)
    if mask.shape[0] != imageShape[0] or mask.shape[1] != imageShape[1]:
        bbox = getBboxFromName(imagePath)
        mask = expand_mask(bbox, mask, image_shape=imageShape)
        mask = mask.astype(np.uint8) * 255
    return mask


def loadOnlyMask(imagePath, imageShape):
    mask = cv2.imread(imagePath, cv2.IMREAD_UNCHANGED)
    if mask.shape[0] != imageShape[0] or mask.shape[1] != imageShape[1]:
        # Finding bbox coordinates from image name
        bbox = getBboxFromName(imagePath)
        shifted = shift_bbox(bbox)
        y1, x1, y2, x2 = shifted

        # Expanding mask to its real size
        mask = expand_mask(shifted, mask, image_shape=shifted[2:])
        mask = mask.astype(np.uint8) * 255
    else:
        # Extracting bbox of
        bbox = extract_bboxes(mask)
        y1, x1, y2, x2 = bbox
    return mask[y1:y2, x1:x2, ...], bbox


def convertImage(inputImagePath: str, outputImagePath: str):
    """
    Convert an image from a format to another one
    :param inputImagePath: path to the initial image
    :param outputImagePath: path to the output image
    :return: None
    """
    image = imread(inputImagePath)
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    cv2.imwrite(outputImagePath, image, CV2_IMWRITE_PARAM)


def getInfoRawDataset(rawDatasetPath: str, verbose=False, adapter: AnnotationAdapter = None, mainFormat="jpg"):
    """
    Listing all available images, those with missing information
    :param verbose: whether or not print should be executed
    :param rawDatasetPath: path to the raw dataset folder
    :param adapter:
    :param mainFormat: the format to use for the dataset
    :return: list of unique files names, list of available images names, list of missing images names,
    list of missing annotations names
    """
    names = []
    images = []  # list of image that can be used to compute masks
    missingImages = []  # list of missing images
    missingAnnotations = []  # list of missing annotations
    inputFormats = ["jpg", "jp2", "png"]
    if adapter is None:
        annotationFormats = adapt.ANNOTATION_FORMAT
    else:
        annotationFormats = [adapter.getAnnotationFormat()]
    fileList = os.listdir(rawDatasetPath)
    if verbose:
        progressBar(0, len(fileList), "Listing files")
    for idx, file in enumerate(fileList):
        if verbose and (idx % 10 == 0 or idx + 1 == len(fileList)):
            lastIdx = idx
            progressBar(idx + 1, len(fileList), "Listing files")
        name = file.split('.')[0]
        if name not in names:  # We want to do this only once per unique file name (without extension)
            names.append(name)

            availableFormat = []
            for format in inputFormats:
                imgPath = os.path.join(rawDatasetPath, f"{name}.{format}")
                if os.path.exists(imgPath):
                    availableFormat.append(format)

            # Testing if annotation file exists for that name
            annotationsExist = False
            for ext in annotationFormats:
                annotationsExist = annotationsExist or os.path.exists(os.path.join(rawDatasetPath, f"{name}.{ext}"))
            if len(availableFormat) > 0:  # At least one image exists
                if not annotationsExist:  # Annotations are missing
                    missingAnnotations.append(name)
                else:
                    if mainFormat not in availableFormat:
                        for format in inputFormats:
                            if format in availableFormat:
                                sourcePath = os.path.join(rawDatasetPath, f"{name}.{format}")
                                outputPath = os.path.join(rawDatasetPath, f"{name}.{mainFormat}")
                                convertImage(sourcePath, outputPath)
                                break
                    images.append(name)  # Adding this image to the list
            elif annotationsExist:  # There is no image file but xml found
                missingImages.append(name)
    if verbose:
        if lastIdx + 1 != len(fileList):
            progressBar(1, 1, "Listing files")
        # Displaying missing image files
        problem = False
        nbMissingImg = len(missingImages)
        if nbMissingImg > 0:
            problem = True
            print('Missing {} image{} : {}'.format(nbMissingImg, 's' if nbMissingImg > 1 else '', missingImages))

        # Displaying missing annotations files
        nbMissingAnnotations = len(missingAnnotations)
        if nbMissingAnnotations > 0:
            problem = True
            print('Missing {} annotation{} : {}'.format(nbMissingAnnotations, 's' if nbMissingAnnotations > 1 else '',
                                                        missingAnnotations))

        # Checking if there is file that is not image nor annotation
        nbImages = len(images)
        if len(names) - nbMissingImg - nbMissingAnnotations - nbImages != 0:
            problem = True
            print('Be careful, there are not only required dataset files in this folder')

        if not problem:
            print("Raw Dataset has no problem. Number of Images : {}".format(nbImages))
    return names, images, missingImages, missingAnnotations


def startWrapper(rawDatasetPath: str, datasetName: str = 'dataset_train', deleteBaseMasks=False,
                 adapter: AnnotationAdapter = None, resize=None, mode="main", classesInfo=None, imageFormat="jpg"):
    """
    Start wrapping the raw dataset into the wanted format
    :param rawDatasetPath: path to the folder containing images and associated annotations
    :param datasetName: name of the output dataset
    :param deleteBaseMasks: delete the base masks images after fusion
    :param adapter: Adapter to use to read annotations, if None compatible adapter will be searched
    :param resize: If tuple given, the images and their masks will be resized to the tuple value
    :param mode: Mode to use
    :param classesInfo: Information about the classes that will be used
    :param imageFormat: the image format to use in the dataset
    :return: None
    """
    names, images, missingImages, missingAnnotations = getInfoRawDataset(rawDatasetPath, verbose=True,
                                                                         adapter=adapter, mainFormat=imageFormat)
    if classesInfo is None:
        if mode == 'main':
            classesInfo = NEPHRO_CLASSES
        elif mode == "cortex":
            classesInfo = CORTICES_CLASSES
        elif mode == "mest_glom":
            classesInfo = MESTC_GLOM_CLASSES
        elif mode == "inflammation":
            classesInfo = INFLAMMATION_CLASSES
    # Creating masks for any image which has all required files and displaying progress
    start_time = time()
    for index, file in enumerate(images):
        progressBar(index, len(images), prefix='Creating masks',
                    suffix=f" {formatTime(round(time() - start_time))} Current : {file}")
        createMasksOfImage(rawDatasetPath, file, datasetName, adapter, classesInfo=classesInfo,
                           resize=resize, imageFormat=imageFormat)
        if mode == "main":
            fuseClassMasks(datasetName, file, "cortex", deleteBaseMasks=deleteBaseMasks, silent=True)
            cleanImage(datasetName, file, cleaningClasses='cortex')
        elif mode == "mest_glom":
            cleanImage(datasetName, file, cleaningClasses='nsg', cleanMasks=True)
        elif mode == "inflammation":
            cleanImage(datasetName, file, cleaningClasses='cortex', cleanMasks=True,
                       excludeClasses=[c["name"] for c in NEPHRO_CLASSES if not c["ignore"]])
    progressBar(1, 1, prefix='Creating masks', suffix=f"{formatTime(round(time() - start_time))}" + " " * 25)

Skinet (Segmentation of the Kidney through a Neural nETwork) Project Dataset tools  Copyright (c) 2021 Skinet Team Licensed under the MIT License (see LICENSE for details) Written by <NAME> Return bbox from a points array     :param pts: the points coordinates     :return: y1, x1, y2, x2 Create the mask image based on its polygon points     :param imgName: name w/o extension of the base image     :param imgShape: shape of the image     :param idMask: the ID of the mask, a number not already used for that image     :param ptsMask: array of [x, y] coordinates which are all the polygon points representing the mask     :param datasetName: name of the output dataset     :param maskClass: name of the associated class of the current mask     :param imageFormat: output format of the masks' images     :param config: config object     :return: None # https://www.programcreek.com/python/example/89415/cv2.fillPoly # Formatting coordinates matrix to get int # Creating black matrix with same size than original image and then drawing the mask # Creating black matrix with same size than original image and then drawing the mask # Saving result image Return the bbox coordinates stored in the image name     :param imageName: the image name from which you want the bbox     :return: the bbox as [y1, x1, y2, x2] Resize mask's base points to fit the targeted size     :param baseMasks array of [x, y] coordinates which are all the polygon points representing the mask     :param xRatio width ratio that will be applied to coordinates     :param yRatio height ratio that will be applied to coordinates Create all the masks of a given image by parsing xml annotations file     :param rawDatasetPath: path to the folder containing images and associated annotations     :param imgName: name w/o extension of an image     :param datasetName: name of the output dataset     :param adapter: the annotation adapter to use to create masks, if None looking for an adapter that can read the file     :param classesInfo: Information about all classes that are used, by default will be nephrology classes Info     :param imageFormat: output format of the image and masks     :param resize: if the image and masks have to be resized     :param config: config object     :return: None # Getting shape of original image (same for all this masks) # Copying the original image in the dataset # TODO use file copy if unchanged else cv2 # Finding annotation files # Choosing the adapter to use (parameters to force it ?) # No adapter given, we are looking for the adapter with highest priority level that can read an/the annotation # file # Using given adapter, we are looking for a file that can be read # Getting the masks data # Creating masks # Converting class id to class name if needed Fuse each targeted-class masks into one     :param datasetPath: the dataset that have been wrapped     :param imageName: the image you want its targeted-class to be fused     :param targetedClass: the class of the masks that have to be fused     :param imageFormat: format to use to save the final targeted-class masks     :param deleteBaseMasks: delete the base masks images after fusion     :param silent: if True will not print     :return: None # Getting the image directory path # Adding each mask to the same image # Saving the fused mask image # Deleting each targeted-class mask except the fused one Cleaning all fused-class directories in the dataset, keeping only unique file or fused ones     :param datasetPath: the dataset that have been wrapped     :param fusedClass: the class      :return: None # Deleting only if strictly more than one fused-class mask image present # Deleting each fused-class mask except the fused one Creating the full_images directory and cleaning the base image by removing non-cleaning-class areas     :param excludeClasses:     :param datasetPath: the dataset that have been wrapped     :param imageName: the image name     :param cleaningClasses: the class to use to clean the image     :param cleanMasks: if true, will clean masks based on the cleaning-class-mask     :param imageFormat: the image format to use to save the image     :param minAreaThreshold: remove mask if its area is smaller than this threshold     :param config: config object     :return: None # Getting the base image # Fusing all the cleaning-class masks and then cleaning the image and if needed the masks # Copying the full image into the correct directory # Cleaning the image and saving it # Cleaning masks so that they cannot exist elsewhere # Checking only for the other classes folder # For each mask of the folder # If mask is not empty # Cleaning it with the cleaning-class masks # If mask was empty or too small after cleaning, we remove it # If mask has is different after cleaning, we replace the original one # If we could not retrieve the original mask ID, give it a unique one # If mini-mask are enabled, we minimize it before saving it # Saving cleaned mask # Remove former mask if not the same name # Adding additional masks # Finding bbox coordinates from image name # Expanding mask to its real size # Extracting bbox of Convert an image from a format to another one     :param inputImagePath: path to the initial image     :param outputImagePath: path to the output image     :return: None Listing all available images, those with missing information     :param verbose: whether or not print should be executed     :param rawDatasetPath: path to the raw dataset folder     :param adapter:     :param mainFormat: the format to use for the dataset     :return: list of unique files names, list of available images names, list of missing images names,     list of missing annotations names # list of image that can be used to compute masks # list of missing images # list of missing annotations # We want to do this only once per unique file name (without extension) # Testing if annotation file exists for that name # At least one image exists # Annotations are missing # Adding this image to the list # There is no image file but xml found # Displaying missing image files # Displaying missing annotations files # Checking if there is file that is not image nor annotation Start wrapping the raw dataset into the wanted format     :param rawDatasetPath: path to the folder containing images and associated annotations     :param datasetName: name of the output dataset     :param deleteBaseMasks: delete the base masks images after fusion     :param adapter: Adapter to use to read annotations, if None compatible adapter will be searched     :param resize: If tuple given, the images and their masks will be resized to the tuple value     :param mode: Mode to use     :param classesInfo: Information about the classes that will be used     :param imageFormat: the image format to use in the dataset     :return: None # Creating masks for any image which has all required files and displaying progress