from ultralytics import YOLO
import os 
import json
from PIL import Image 
from uuid import uuid4 as uuid

# Load Model
model = YOLO('./model/0307241454.pt')

batchNumber = 18
scanningDir = f'C:\\Users\\epicl\\Desktop\\thoth Code\\label-studio\\images\\batch{batchNumber}'
outputDir = f'C:\\Users\\epicl\\Desktop\\thoth Code\\label-studio\\images\\batch{batchNumber}Output'
if not os.path.exists(outputDir):
    os.makedirs(outputDir)
if not os.path.exists(scanningDir):
    print('Wrong Path')
    Exception('Wrong Path')

imageNames = os.listdir(scanningDir)
imageNames = [x for x in imageNames if x.endswith('.jpg') or x.endswith('.png') or x.endswith('.jpeg') or x.endswith('.JPG') or x.endswith('.PNG') or x.endswith('.JPEG')]
imageNames = [os.path.join(scanningDir, x) for x in imageNames]

# Get Results in Batches
results = []
batchSize = 32
number_of_batches = len(imageNames) // batchSize
if len(imageNames) % batchSize != 0:
    number_of_batches += 1
for i in range(number_of_batches):
    start = i*batchSize
    end = (i+1)*batchSize
    if end > len(imageNames):
        end = len(imageNames)
    batch = imageNames[start:end]
    batchClean = []
    for imageFilename in batch:
        try:
            image = Image.open(imageFilename)
            batchClean.append(imageFilename)
        except:
            print(f'Error with {imageFilename}')
    results += model(batchClean)
    print(f'Batch {start} to {end} done')
class_names = model.names


# Save as individual then batch import 

for idx, result in enumerate(results):
    # Get Image Dimensions
    image = Image.open(imageNames[idx])
    image_width, image_height = image.size

    # Sort out the keys annotations first
    bboxes = result.cpu().to('cpu').numpy().boxes.data
    local_dict = {x:[] for x in class_names.values()}

    for bbox in bboxes:
        x1,y1,x2,y2, confidence , classIdx = bbox[0], bbox[1], bbox[2], bbox[3], bbox[4], bbox[5]
        x = x1
        y = y1
        w = x2-x1
        h = y2-y1
        classname = class_names[classIdx]
        
        local_dict[classname].append([float(x), float(y), float(w), float(h), float(confidence)])

    # clean 0 lengths Labels
    localCopy= local_dict.copy()
    for key in localCopy.keys():
        if len(local_dict[key]) == 0:
            del local_dict[key]

    formatJson = [
        {
            "data":{
                'image': f'http://localhost:8080/data/local-files/?d=batch{batchNumber}/{os.path.basename(imageNames[idx])}'
            },
            "predictions": [
                {
                    "result": [
                    ]
                }
            ]
        }
    ]

    for className in local_dict.keys():
        for bbox in local_dict[className]:
            x,y,w,h, confidence = bbox
            percentage_x = x / image_width * 100
            percentage_y = y / image_height * 100 
            percentage_w = w / image_width * 100
            percentage_h = h / image_height * 100
            outputJsonLocal ={
                            "id": str(uuid()),
                            "type": "rectanglelabels", 
                            "from_name": "label",
                            "to_name": "image",
                            "original_width": image_width, 
                            "original_height": image_height,
                            "image_rotation": 0,
                            "origin": "manual",
                            "value": {
                                "rotation": 0,
                                "x": percentage_x, #bounding box x
                                "y": percentage_y, #bounding box y
                                "width": percentage_w, #bounding box width
                                "height": percentage_h, #bounding box height
                                "rotation" : 0,
                                "rectanglelabels": [className] #class name
                            }

                        }
            formatJson[0]['predictions'][0]['result'].append(outputJsonLocal)
    
    save_name = os.path.basename(imageNames[idx]).split('.')[0] + '.json'
    
    with open (os.path.join(outputDir,save_name ), 'w') as f:
        f.write(json.dumps(formatJson, indent=4))