Fuururuny's picture
Update app.py
bedc95f verified
Raw
History Blame Contribute Delete
1.42 kB
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()