app.py

import os
import gc
import json
import ast
import re
import base64
from io import BytesIO
from threading import Thread
from typing import Tuple

import gradio as gr
from gradio import Server
from fastapi.responses import HTMLResponse
import spaces
import torch
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import supervision as sv
from transformers import (
    Qwen3_5ForConditionalGeneration,
    AutoProcessor,
    TextIteratorStreamer,
)

# ------------------------------------------------------------------
# Config & Constants
# ------------------------------------------------------------------
MODEL_NAME = "Qwen/Qwen3.8-27B"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16

GPU_DURATIONS = [60, 90, 120, 150, 180, 250, 300]
DEFAULT_GPU_DURATION_IDX = 1

BRIGHT_YELLOW = sv.Color(r=255, g=230, b=0)
DARK_OUTLINE = sv.Color(r=40, g=40, b=40)
BLACK = sv.Color(r=0, g=0, b=0)
WHITE = sv.Color(r=255, g=255, b=255)

SPATIAL_LINE = (255, 69, 0)       # OrangeRed
SPATIAL_DOT = (255, 69, 0)
SPATIAL_RING = (255, 255, 255)
SPATIAL_LABEL_BG = (50, 10, 0)
SPATIAL_LABEL_TXT = (255, 255, 255)
SPATIAL_ARROW = (255, 140, 0)

EXAMPLES_CONFIG = [
    {"image": "examples/1.jpg", "prompt": "Detect the yellow car that is parked.", "task": "Detect"},
    {"image": "examples/4.jpg", "prompt": "Map the waypoints from the basketball to the basket to shoot the ball inside the net.", "task": "Spatial"},   
    {"image": "examples/2.jpg", "prompt": "Point to all the red cars.", "task": "Point"},
    {"image": "examples/3.jpg", "prompt": "Map a path from the door to the lamp.", "task": "Spatial"},
]

# ------------------------------------------------------------------
# Model Loading
# ------------------------------------------------------------------
print(f"Loading model: {MODEL_NAME} ...")
qwen_model = Qwen3_5ForConditionalGeneration.from_pretrained(
    MODEL_NAME, torch_dtype=DTYPE, device_map=DEVICE, attn_implementation="kernels-community/flash-attn2@v3",
).eval()
qwen_processor = AutoProcessor.from_pretrained(MODEL_NAME)
print("Model loaded.")

# ------------------------------------------------------------------
# Helper Functions
# ------------------------------------------------------------------
def make_thumb_b64(path, max_dim=220):
    if not os.path.exists(path):
        return ""
    try:
        img = Image.open(path).convert("RGB")
        img.thumbnail((max_dim, max_dim), Image.LANCZOS)
        buf = BytesIO()
        img.save(buf, format="JPEG", quality=65)
        return f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()).decode()}"
    except Exception:
        return ""

def encode_full_image(path):
    if not os.path.exists(path):
        return ""
    try:
        with open(path, "rb") as f:
            data = f.read()
        ext = path.rsplit(".", 1)[-1].lower()
        mime = {"jpg": "image/jpeg", "jpeg": "image/jpeg", "png": "image/png", "webp": "image/webp"}.get(ext, "image/jpeg")
        return f"data:{mime};base64,{base64.b64encode(data).decode()}"
    except Exception:
        return ""

def build_client_config():
    examples = []
    for i, ex in enumerate(EXAMPLES_CONFIG):
        examples.append({
            "idx": i,
            "thumb": make_thumb_b64(ex["image"]),
            "prompt": ex["prompt"],
            "task": ex["task"],
        })
    return {"examples": examples, "gpu_durations": GPU_DURATIONS, "default_gpu_idx": DEFAULT_GPU_DURATION_IDX}

CLIENT_CONFIG = build_client_config()

def safe_parse_json(text: str):
    text = re.sub(r"```(json)?", "", text).strip()
    match = re.search(r'(\[.*\]|\{.*\})', text, re.DOTALL)
    if match:
        json_str = match.group(1)
        json_str_clean = re.sub(r',\s*([}\]])', r'\1', json_str)
        try:
            return json.loads(json_str_clean)
        except json.JSONDecodeError:
            try:
                return ast.literal_eval(json_str_clean)
            except Exception:
                pass
    text_clean = re.sub(r',\s*([}\]])', r'\1', text)
    try:
        return json.loads(text_clean)
    except json.JSONDecodeError:
        pass
    try:
        return ast.literal_eval(text_clean)
    except Exception:
        pass
    return []

def _extract_point(item: dict):
    if not isinstance(item, dict):
        return None
    for k in ["point_2d", "point", "points", "coordinate", "coordinates", "xy"]:
        if k in item and isinstance(item[k], (list, tuple)) and len(item[k]) == 2:
            return item[k]
    return None

def _extract_bbox(item: dict):
    if not isinstance(item, dict):
        return None
    for k in ["bbox_2d", "bbox", "box", "bounding_box", "xyxy"]:
        if k in item and isinstance(item[k], (list, tuple)) and len(item[k]) == 4:
            return item[k]
    return None

def _load_font(size: int = 16):
    size = max(6, int(size))
    try:
        return ImageFont.truetype("arial.ttf", size)
    except (IOError, OSError):
        try:
            return ImageFont.truetype("DejaVuSans.ttf", size)
        except (IOError, OSError):
            return ImageFont.load_default()

def pil_to_b64_png(image: Image.Image) -> str:
    buf = BytesIO()
    image.save(buf, format="PNG")
    return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"

def annotate_image(image: Image.Image, result: dict, point_radius: int = 6, box_thickness: int = 2, text_scale: float = 0.5):
    if not isinstance(image, Image.Image) or not isinstance(result, dict):
        return image
    image = image.convert("RGB")
    ow, oh = image.size

    point_radius = max(1, int(point_radius))
    box_thickness = max(1, int(box_thickness))
    text_scale = max(0.1, float(text_scale))
    text_thickness = max(1, round(text_scale * 2))

    if "points" in result and result["points"]:
        pts = [[int(p["x"] * ow), int(p["y"] * oh)] for p in result["points"]]
        if not pts:
            return image
        kp = sv.KeyPoints(xy=np.array(pts).reshape(1, -1, 2))
        scene = np.array(image.copy())
        scene = sv.VertexAnnotator(radius=point_radius + 3, color=DARK_OUTLINE).annotate(scene=scene, key_points=kp)
        scene = sv.VertexAnnotator(radius=point_radius, color=BRIGHT_YELLOW).annotate(scene=scene, key_points=kp)
        labels = [p.get("label", "") for p in result["points"]]
        if any(labels):
            tb, vl = [], []
            for i, p in enumerate(result["points"]):
                if labels[i]:
                    cx, cy = int(p["x"] * ow), int(p["y"] * oh)
                    tb.append([cx - 2, cy - 2, cx + 2, cy + 2])
                    vl.append(labels[i])
            if tb:
                scene = sv.LabelAnnotator(
                    color=BRIGHT_YELLOW, text_color=BLACK, text_scale=text_scale,
                    text_thickness=text_thickness, text_padding=5,
                    text_position=sv.Position.TOP_CENTER,
                    color_lookup=sv.ColorLookup.INDEX,
                ).annotate(scene=scene, detections=sv.Detections(xyxy=np.array(tb)), labels=vl)
        return Image.fromarray(scene)

    if "objects" in result and result["objects"]:
        boxes, labels = [], []
        for obj in result["objects"]:
            boxes.append([
                obj.get("x_min", 0.0) * ow, obj.get("y_min", 0.0) * oh,
                obj.get("x_max", 0.0) * ow, obj.get("y_max", 0.0) * oh,
            ])
            labels.append(obj.get("label", "object"))
        if not boxes:
            return image
        scene = np.array(image.copy())
        h, w = scene.shape[:2]
        masks = np.zeros((len(boxes), h, w), dtype=bool)
        for i, box in enumerate(boxes):
            x1, y1 = max(0, int(box[0])), max(0, int(box[1]))
            x2, y2 = min(w, int(box[2])), min(h, int(box[3]))
            masks[i, y1:y2, x1:x2] = True
        dets = sv.Detections(xyxy=np.array(boxes), mask=masks)
        if len(dets) == 0:
            return image
        scene = sv.MaskAnnotator(color=BRIGHT_YELLOW, opacity=0.18, color_lookup=sv.ColorLookup.INDEX).annotate(scene=scene, detections=dets)
        scene = sv.BoxAnnotator(color=BRIGHT_YELLOW, thickness=box_thickness, color_lookup=sv.ColorLookup.INDEX).annotate(scene=scene, detections=dets)
        scene = sv.LabelAnnotator(
            color=BRIGHT_YELLOW, text_color=BLACK, text_scale=text_scale,
            text_thickness=text_thickness, text_padding=6, color_lookup=sv.ColorLookup.INDEX,
        ).annotate(scene=scene, detections=dets, labels=labels)
        return Image.fromarray(scene)
    return image

def annotate_spatial_path(image: Image.Image, result: dict, dot_radius: int = 6, line_width: int = 4, text_scale: float = 0.5):
    if not isinstance(image, Image.Image) or not isinstance(result, dict):
        return image
    image = image.convert("RGB")
    w, h = image.size
    if "points" not in result or not result["points"]:
        return image

    dot_radius = max(1, int(dot_radius))
    line_width = max(1, int(line_width))
    text_scale = max(0.1, float(text_scale))

    draw = ImageDraw.Draw(image, "RGBA")
    font_label = _load_font(16 * text_scale * 2)
    font_num = _load_font(14 * text_scale * 2)

    points = result["points"]
    pts = [(int(p["x"] * w), int(p["y"] * h)) for p in points]
    labels = [p.get("label", f"P{i+1}") for i, p in enumerate(points)]

    scale_ratio = dot_radius / 8.0

    if len(pts) >= 2:
        for i in range(len(pts) - 1):
            draw.line([pts[i], pts[i+1]], fill=SPATIAL_LINE + (60,), width=line_width + 6)
        for i in range(len(pts) - 1):
            draw.line([pts[i], pts[i+1]], fill=SPATIAL_LINE, width=line_width)

        for i in range(len(pts) - 1):
            x1, y1 = pts[i]
            x2, y2 = pts[i+1]
            dx, dy = x2 - x1, y2 - y1
            length = (dx * dx + dy * dy) ** 0.5
            if length < 12:
                continue
            ux, uy = dx / length, dy / length
            offset = 18 * scale_ratio
            bx, by = x2 - ux * offset, y2 - uy * offset
            px, py = -uy, ux
            aw, ah = 7 * scale_ratio, 9 * scale_ratio
            p1 = (bx + px * aw, by + py * aw)
            p2 = (bx - px * aw, by - py * aw)
            p3 = (bx + ux * ah, by + uy * ah)
            draw.polygon([p1, p2, p3], fill=SPATIAL_ARROW)

    for i, (cx, cy) in enumerate(pts):
        halo_r = dot_radius + 8
        ring_r = dot_radius + 3
        draw.ellipse((cx - halo_r, cy - halo_r, cx + halo_r, cy + halo_r), fill=SPATIAL_LINE + (50,))
        draw.ellipse((cx - ring_r, cy - ring_r, cx + ring_r, cy + ring_r),
                     outline=SPATIAL_RING, width=max(1, round(3 * scale_ratio)))
        draw.ellipse((cx - dot_radius, cy - dot_radius, cx + dot_radius, cy + dot_radius),
                     fill=SPATIAL_DOT, outline=SPATIAL_DOT)
        num_text = str(i + 1)
        nbbox = draw.textbbox((0, 0), num_text, font=font_num)
        nw = nbbox[2] - nbbox[0]
        nh = nbbox[3] - nbbox[1]
        nx, ny = cx - nw // 2, cy - nh // 2 - 1
        draw.text((nx, ny), num_text, fill=SPATIAL_RING, font=font_num)

    for i, (cx, cy) in enumerate(pts):
        label = f"{i+1}. {labels[i]}"
        bbox = draw.textbbox((0, 0), label, font=font_label)
        tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
        lx, ly = cx + dot_radius + 10, cy - th - 8
        pad = 5
        draw.rectangle(
            (lx - pad, ly - pad, lx + tw + pad, ly + th + pad),
            fill=SPATIAL_LABEL_BG,
            outline=SPATIAL_LINE, width=1,
        )
        draw.text((lx, ly), label, fill=SPATIAL_LABEL_TXT, font=font_label)

    n_pts = len(pts)
    legend_text = f"Spatial map · {n_pts} waypoints · path length {len(pts)-1} segments"
    legend_font = _load_font(13 * text_scale * 2)
    bbox = draw.textbbox((0, 0), legend_text, font=legend_font)
    tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
    fx, fy = 10, h - th - 22
    draw.rectangle((fx - 8, fy - 6, fx + tw + 16, fy + th + 10),
                   fill=SPATIAL_LABEL_BG + (220,))
    draw.text((fx, fy), legend_text, fill=SPATIAL_LABEL_TXT, font=legend_font)

    return image

def get_gpu_duration(image_b64, prompt, task_type, point_radius, box_thickness, text_scale, gpu_duration_seconds):
    try:
        return int(gpu_duration_seconds)
    except (TypeError, ValueError):
        return GPU_DURATIONS[DEFAULT_GPU_DURATION_IDX]

# ------------------------------------------------------------------
# Gradio Server (Server mode)
# ------------------------------------------------------------------
app = Server(title="Qwen3.8-27B-Object-Detection")

@app.mcp.tool(name="run_inference")
@app.api(name="run_inference")
@spaces.GPU(size="xlarge", duration=get_gpu_duration)
def infer(
    image_b64: str,
    prompt: str,
    task_type: str,
    point_radius: int,
    box_thickness: int,
    text_scale: float,
    gpu_duration_seconds: int,
) -> dict:
    """Runs object detection, point localization, or spatial mapping."""
    gc.collect()
    torch.cuda.empty_cache()

    if not image_b64:
        raise gr.Error("Please upload an image.")
    if not prompt or prompt.strip() == "":
        raise gr.Error("Please provide a prompt.")

    try:
        header, data = image_b64.split(",", 1)
        pil_image = Image.open(BytesIO(base64.b64decode(data))).convert("RGB")
    except Exception as e:
        raise gr.Error(f"Invalid image data: {e}")

    pil_image.thumbnail((512, 512))

    if task_type == "Detect":
        full_prompt = (
            f"Provide bounding box coordinates for {prompt}. "
            f"Report strictly in JSON format as a list of objects with 'label' and "
            f"'bbox_2d' (xmin, ymin, xmax, ymax in 0-1000 scale)."
        )
    elif task_type == "Point":
        full_prompt = (
            f"Provide 2d point coordinates for {prompt}. "
            f"Report strictly in JSON format as a list of objects with 'label' and "
            f"'point_2d' (x, y in 0-1000 scale)."
        )
    elif task_type == "Spatial":
        full_prompt = (
            f"Identify the key spatial waypoints to map a path/route for: {prompt}. "
            f"Return the points in the order they should be connected along the path, "
            f"from start to end. Report in JSON format as a list of objects with "
            f"'label' (name of each waypoint) and 'point_2d' (x, y in 0-1000 scale). "
            f"Include 2-8 waypoints that best describe the spatial map."
        )
    else:
        full_prompt = prompt

    messages = [{
        "role": "user",
        "content": [
            {"type": "image", "image": pil_image},
            {"type": "text", "text": full_prompt},
        ],
    }]
    text = qwen_processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = qwen_processor(
        text=[text], images=[pil_image],
        return_tensors="pt", padding=True,
    ).to(qwen_model.device)

    streamer = TextIteratorStreamer(
        qwen_processor.tokenizer,
        skip_prompt=True, skip_special_tokens=True, timeout=120,
    )
    thread = Thread(
        target=qwen_model.generate,
        kwargs=dict(
            **inputs,
            streamer=streamer,
            max_new_tokens=2048,
            use_cache=True,
            do_sample=False,
        ),
    )
    thread.start()
    full_text = ""
    for tok in streamer:
        full_text += tok
    thread.join()

    if task_type == "Point":
        parsed = safe_parse_json(full_text)
        if isinstance(parsed, dict):
            list_found = False
            for k in ["points", "keypoints", "point"]:
                if k in parsed and isinstance(parsed[k], list):
                    parsed = parsed[k]
                    list_found = True
                    break
            if not list_found:
                for v in parsed.values():
                    if isinstance(v, list):
                        parsed = v
                        break
                else:
                    parsed = []

        result = {"points": []}
        if isinstance(parsed, list):
            for item in parsed:
                pt = _extract_point(item)
                if pt:
                    x, y = pt
                    result["points"].append({
                        "label": item.get("label", ""),
                        "x": x / 1000.0, "y": y / 1000.0,
                    })

        if result["points"]:
            annotated_img = annotate_image(
                pil_image.copy(), result,
                point_radius=point_radius, box_thickness=box_thickness, text_scale=text_scale,
            )
            return {"image": pil_to_b64_png(annotated_img), "text": json.dumps(result, indent=2)}
        else:
            return {"image": pil_to_b64_png(pil_image), "text": f"Could not extract any points.\nRaw model output:\n{full_text}"}

    elif task_type == "Detect":
        parsed = safe_parse_json(full_text)
        if isinstance(parsed, dict):
            list_found = False
            for k in ["objects", "detections", "bboxes", "boxes", "results"]:
                if k in parsed and isinstance(parsed[k], list):
                    parsed = parsed[k]
                    list_found = True
                    break
            if not list_found:
                for v in parsed.values():
                    if isinstance(v, list):
                        parsed = v
                        break
                else:
                    parsed = []

        result = {"objects": []}
        if isinstance(parsed, list):
            for item in parsed:
                bbox = _extract_bbox(item)
                if bbox:
                    xmin, ymin, xmax, ymax = bbox
                    result["objects"].append({
                        "label": item.get("label", "object"),
                        "x_min": xmin / 1000.0, "y_min": ymin / 1000.0,
                        "x_max": xmax / 1000.0, "y_max": ymax / 1000.0,
                    })

        if result["objects"]:
            annotated_img = annotate_image(
                pil_image.copy(), result,
                point_radius=point_radius, box_thickness=box_thickness, text_scale=text_scale,
            )
            return {"image": pil_to_b64_png(annotated_img), "text": json.dumps(result, indent=2)}
        else:
            return {"image": pil_to_b64_png(pil_image), "text": f"Could not extract any objects.\nRaw model output:\n{full_text}"}

    elif task_type == "Spatial":
        parsed = safe_parse_json(full_text)
        if isinstance(parsed, dict):
            list_found = False
            for k in ["points", "waypoints", "path", "route", "nodes", "map"]:
                if k in parsed and isinstance(parsed[k], list):
                    parsed = parsed[k]
                    list_found = True
                    break
            if not list_found:
                for v in parsed.values():
                    if isinstance(v, list):
                        parsed = v
                        break
                else:
                    parsed = []

        result = {"points": []}
        if isinstance(parsed, list):
            for item in parsed:
                pt = _extract_point(item)
                if pt:
                    x, y = pt
                    result["points"].append({
                        "label": item.get("label", "waypoint"),
                        "x": x / 1000.0, "y": y / 1000.0,
                    })

        if result["points"]:
            wp_lines = "\n".join(
                f"  {i+1}. {p['label']}  →  ({p['x']:.3f}, {p['y']:.3f})"
                for i, p in enumerate(result["points"])
            )
            summary = (
                f"Spatial map generated.\n"
                f"Waypoints ({len(result['points'])}):\n{wp_lines}\n"
                f"Path segments: {max(0, len(result['points']) - 1)}"
            )
            annotated_img = annotate_spatial_path(
                pil_image.copy(), result,
                dot_radius=point_radius, line_width=box_thickness * 2, text_scale=text_scale,
            )
            return {"image": pil_to_b64_png(annotated_img), "text": summary}
        else:
            return {"image": pil_to_b64_png(pil_image), "text": f"Could not extract any spatial waypoints.\nRaw model output:\n{full_text}"}

    return {"image": pil_to_b64_png(pil_image), "text": "Unknown task type."}


@app.api(name="load_example", queue=False)
def load_example(idx: float) -> dict:
    try:
        i = int(idx)
    except (ValueError, TypeError):
        i = -1
    if i < 0 or i >= len(EXAMPLES_CONFIG):
        return {"image": "", "prompt": "", "task": "", "status": "error"}
    ex = EXAMPLES_CONFIG[i]
    b64 = encode_full_image(ex["image"])
    return {
        "image": b64,
        "prompt": ex["prompt"],
        "task": ex["task"],
        "status": "ok" if b64 else "error"
    }

@app.get("/api/config")
def client_config():
    return CLIENT_CONFIG

@app.get("/", response_class=HTMLResponse)
async def homepage():
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return f.read()

if __name__ == "__main__":
    app.launch(show_error=True, mcp_server=True)