from helpers.AnnotationLocationLoader import AnnotationLocationLoader
from helpers.PredictionLocationLoader import PredictionLocationLoader
from pathlib import Path
from optparse import OptionParser
import shutil

def convertGroundtruths(out_path):
    # create folder for gt
    out_path = out_path / 'groundtruths'
    if not out_path.exists():
        out_path.mkdir(parents=True)

    annotLoader = AnnotationLocationLoader(annotation_file='../input/caries_dataset_annotation.json', images_base_folder='../input/evaluation_data')
    for image in annotLoader.get_all_available_images():
        #create file
        file = image[:-3] + 'txt'
        with open(out_path  / file, "w") as groundtruth_file:
            for annotation in annotLoader.get_annotations(image):
                # format: <class_name> <left> <top> <right> <bottom>
                if not str(annotation[0]) == 'caries':
                    continue
                groundtruth_file.write(str(annotation[0]) + " " + str(annotation[1][0][0]) + " " + str(annotation[1][0][1]) + " " + str(annotation[1][1][0]) + " " + str(annotation[1][1][1]) + "\n")

def convertDetections(out_path, predictions_file):
    # create folder for detections
    out_path = out_path / 'detections'
    if not out_path.exists():
        out_path.mkdir(parents=True)

    predLoader = PredictionLocationLoader(prediction_file=predictions_file, images_base_folder='../input/evaluation_data')
    for image in predLoader.get_all_annotated_images():
        file = image[:-3] + 'txt'
        with open(out_path / file, "w") as detection_file:
            for annotation in predLoader.get_annotations(image):
                # format: <class_name> <confidence> <left> <top> <right> <bottom>
                detection_file.write(
                    str(annotation[0]) + " " + str(annotation[1][2]) + " " + str(annotation[1][0][0]) + " " + str(annotation[1][0][1]) + " " + str(
                        annotation[1][1][0]) + " " + str(annotation[1][1][1]) + "\n")


'''
    Please run the Pascal VOC Metrics implemented by Rafael Padilla on the output of this Script.
    https://github.com/rafaelpadilla/Object-Detection-Metrics
'''

parser = OptionParser()

parser.add_option("-i", "--predictions_file", dest="predictions_file", help="Path to predictions file as input.", default="../out/evaluation_ws_base/predictions.txt")
parser.add_option("-o", "--path_output", dest="output_path", help="Path to output folder.", default='../out/convertedForObjectDetectionMetrics_ws_base')

(options, args) = parser.parse_args()

if not Path(options.predictions_file).is_file():
    parser.error('Error: Could not find the prediction file. Check the path passed to -i or --prediction_file')

output_path = Path(options.output_path)
shutil.rmtree(str(output_path))
predictions_file = options.predictions_file

convertGroundtruths(output_path)
convertDetections(output_path, predictions_file)

print('[INFO] Finished conversion for ObjectDetectionMetrics.')
print('[INFO] Saved to folder {}'.format(str(output_path)))