brasilin2006/DINO-SDK-FASTAPI

1import io, base64, re, os
2from typing import List
3from PIL import Image
4
5import torch
6from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
7import gradio as gr
8
9MODEL_ID = os.environ.get("MODEL_ID", "IDEA-Research/grounding-dino-tiny")
10DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
11
12processor = AutoProcessor.from_pretrained(MODEL_ID)
13model = AutoModelForZeroShotObjectDetection.from_pretrained(MODEL_ID).to(DEVICE).eval()
14
15data_url_re = re.compile(r"^data:image/(png|jpeg|jpg);base64,(.+)$", re.I)
16
17def decode_image_from_data_url(data_url: str) -> Image.Image:
18    m = data_url_re.match(data_url)
19    if not m:
20        raise ValueError("Invalid image data URL")
21    img_bytes = base64.b64decode(m.group(2))
22    return Image.open(io.BytesIO(img_bytes)).convert("RGB")
23
24def detect_core(image: Image.Image, prompts: List[str], threshold: float):
25    W, H = image.size
26    text = ". ".join(s.strip() for s in prompts if s.strip())
27    if not text.endswith("."):
28        text += "."
29    inputs = processor(images=image, text=[[text]], return_tensors="pt").to(DEVICE)
30    with torch.no_grad():
31        outputs = model(**inputs)
32    results = processor.post_process_grounded_object_detection(
33        outputs,
34        inputs.input_ids,
35        box_threshold=max(0.05, min(0.99, threshold)),
36        text_threshold=0.25,
37        target_sizes=[(H, W)]
38    )[0]
39    dets = []
40    for i, (xyxy, score, label) in enumerate(zip(results["boxes"], results["scores"], results["labels"])):
41        x1, y1, x2, y2 = [float(v) for v in xyxy.tolist()]
42        w, h = (x2 - x1), (y2 - y1)
43        if w <= 0 or h <= 0:
44            continue
45        dets.append({
46            "id": f"d{i}",
47            "label": str(label),
48            "confidence": float(score.item()),
49            "bbox": [x1, y1, w, h],
50            "angle": 0.0
51        })
52    return dets
53
54def detect_api(data_url: str, prompts_csv: str, threshold: float):
55    # Accept your front end’s Data URL verbatim
56    image = decode_image_from_data_url(data_url)
57    prompts = [s.strip() for s in prompts_csv.split(",") if s.strip()]
58    return detect_core(image, prompts, threshold)
59
60with gr.Blocks() as demo:
61    gr.Markdown("### Grounding DINO detect (named API)")
62    with gr.Row():
63        with gr.Column():
64            data_url = gr.Textbox(label="Image Data URL (data:image/...;base64,...)")
65            prm = gr.Textbox(value="person, car, tree, bench", label="Prompts (comma-separated)")
66            thr = gr.Slider(0, 1, value=0.25, step=0.01, label="Threshold")
67            btn = gr.Button("Run detect")
68        out = gr.JSON(label="Detections")
69    btn.click(detect_api, inputs=[data_url, prm, thr], outputs=[out], api_name="detect")
70
71demo.queue(concurrency_count=1).launch()
72