| 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() |
