import subprocess
from ultralytics import YOLO
import cv2
import numpy as np
import gradio as gr
from io import BytesIO
from cairosvg import svg2png
import datetime
#===== load model =====
model = YOLO('lite.pt')

#===== some funcs =====
def prepare_image(svg):
        data = BytesIO(svg2png(svg))
        image = cv2.imdecode(np.frombuffer(data.read(), np.uint8), cv2.IMREAD_COLOR)
        return image
def processing(image):
    results = model(image, conf=0.5, max_det=4)[0]
    print(datetime.datetime.now())
    
    classes_names = results.names
    classes = results.boxes.cls.cpu().numpy()
    boxes = results.boxes.xyxy.cpu().numpy().astype(np.int32)

    objects = []

    for class_id, box in zip(classes, boxes):
        class_name = classes_names[int(class_id)]

        x1, y1, x2, y2 = box
        objects.append((class_name, x1))

    objects.sort(key=lambda obj: obj[1])
    text_result = ''.join([obj[0] for obj in objects])
    return text_result
def predict(svgdata):
    image = prepare_image(svgdata)
    if image.size > 630000:
        print("bad image")
        return;
    return processing(image)

title = "Made with love❤️❤️❤️❤️\nby nof"
description = "bruh4"

iface = gr.Interface(fn=predict, 
                     inputs=gr.Textbox(), 
                     outputs=gr.Textbox(),
                     title=title,
                     description=description)
iface.launch()