import gradio as gr
import io, json, base64
from PIL import Image
from ultralytics import YOLO

model = YOLO("comic_bubble_detector.pt")

def detect(image_b64: str, conf: float = 0.3):
    img_bytes = base64.b64decode(image_b64)
    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    results = model(img, conf=conf, verbose=False)[0]
    print(f"🔍 كشف {len(results.boxes)} بالون بـ conf={conf}", flush=True)
    labels_map = {0: "bubble", 1: "text_bubble", 2: "text_free"}
    bubbles = []
    for box in results.boxes:
        cls = int(box.cls[0])
        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
        crop = img.crop((x1, y1, x2, y2))
        buf = io.BytesIO()
        crop.save(buf, format="JPEG", quality=95)
        crop_b64 = base64.b64encode(buf.getvalue()).decode()
        bubbles.append({
            "box": [x1, y1, x2, y2],
            "class": cls,
            "label": labels_map.get(cls, "unknown"),
            "conf": float(box.conf[0]),
            "crop_b64": crop_b64
        })
    return json.dumps(bubbles)

gr.Interface(fn=detect, inputs=["text", "number"], outputs="text", api_name="detect").launch()