app.py

import os
import json
import math
import tempfile
from typing import Any, Dict, List, Tuple

import gradio as gr
import numpy as np
import cv2
from PIL import Image, ImageDraw, ImageFont

from inference_sdk import InferenceHTTPClient
from inference_sdk.http.errors import HTTPCallErrorError


# ---- Roboflow client ----
client = InferenceHTTPClient(
    api_url="https://serverless.roboflow.com",
    api_key=os.environ["ROBOFLOW_API_KEY"]
)

WORKSPACE = "ata-assignment-1-mkqz4"
WORKFLOW_ID = "custom-workflow-4"


# ---- Helpers: parse + draw ----
def _extract_predictions(result: Any) -> List[Dict[str, Any]]:
    """
    Workflows can return different shapes. We try common patterns.
    Returns a list of prediction dicts with either:
      - x,y,width,height (+ class, confidence), OR
      - bbox dict
    """
    # Often: result is a list with one item
    item = result[0] if isinstance(result, list) and result else result

    if not isinstance(item, dict):
        return []

    # Common: {"predictions": {"image": {...}, "predictions": [...]}}
    if isinstance(item.get("predictions"), dict) and isinstance(item["predictions"].get("predictions"), list):
        return item["predictions"]["predictions"]

    # Sometimes predictions sit directly under a key
    for v in item.values():
        if isinstance(v, dict) and isinstance(v.get("predictions"), list):
            return v["predictions"]

    # Or result itself contains predictions list
    if isinstance(item.get("predictions"), list):
        return item["predictions"]

    return []


def _draw_boxes_pil(img: Image.Image, preds: List[Dict[str, Any]], conf_thresh: float) -> Image.Image:
    img = img.convert("RGB")
    draw = ImageDraw.Draw(img)

    try:
        font = ImageFont.load_default()
    except Exception:
        font = None

    for p in preds:
        conf = float(p.get("confidence", p.get("conf", 0.0)))
        if conf < conf_thresh:
            continue

        cls = p.get("class", p.get("label", "obj"))

        # xywh center-based
        if all(k in p for k in ["x", "y", "width", "height"]):
            x, y, w, h = float(p["x"]), float(p["y"]), float(p["width"]), float(p["height"])
            x1, y1, x2, y2 = x - w / 2, y - h / 2, x + w / 2, y + h / 2

        # bbox dict
        elif isinstance(p.get("bbox"), dict):
            b = p["bbox"]
            if all(k in b for k in ["x1", "y1", "x2", "y2"]):
                x1, y1, x2, y2 = map(float, (b["x1"], b["y1"], b["x2"], b["y2"]))
            elif all(k in b for k in ["left", "top", "right", "bottom"]):
                x1, y1, x2, y2 = map(float, (b["left"], b["top"], b["right"], b["bottom"]))
            else:
                continue
        else:
            continue

        draw.rectangle([x1, y1, x2, y2], width=3)
        label = f"{cls} {conf:.2f}"
        draw.text((x1, max(0, y1 - 14)), label, font=font)

    return img


def _run_on_image_path(image_path: str, use_cache: bool) -> Any:
    return client.run_workflow(
        workspace_name=WORKSPACE,
        workflow_id=WORKFLOW_ID,
        images={"image": image_path},
        use_cache=use_cache
    )


# ---- Main: Image ----
def infer_image(image_path: str, use_cache: bool, conf_thresh: float):
    if image_path is None:
        return None, {"error": "No image uploaded."}

    try:
        result = _run_on_image_path(image_path, use_cache=use_cache)
        preds = _extract_predictions(result)

        img = Image.open(image_path)
        annotated = _draw_boxes_pil(img, preds, conf_thresh=conf_thresh)
        return annotated, result

    except HTTPCallErrorError as e:
        return None, {
            "error": "Roboflow request failed",
            "status_code": getattr(e, "status_code", None),
            "api_message": getattr(e, "api_message", str(e)),
            "description": str(e),
        }


# ---- Main: Video (frame sampling) ----
def infer_video(video_path: str, use_cache: bool, conf_thresh: float, fps_out: int, sample_every_n: int):
    """
    Reads video, runs workflow on every Nth frame, draws boxes, writes annotated mp4.
    For non-sampled frames, we reuse the last predictions (so boxes persist smoothly).
    """
    if video_path is None:
        return None, {"error": "No video uploaded."}

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None, {"error": "Could not open video."}

    # Video properties
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    in_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0

    # We’ll write a new mp4
    tmp_out = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    tmp_out.close()

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(tmp_out.name, fourcc, float(fps_out), (width, height))

    frame_idx = 0
    last_preds: List[Dict[str, Any]] = []
    summary = {
        "input_fps": in_fps,
        "output_fps": fps_out,
        "sample_every_n_frames": sample_every_n,
        "frames_processed": 0,
        "workflow_calls": 0,
        "example_results": []
    }

    try:
        while True:
            ok, frame = cap.read()
            if not ok:
                break

            # Decide whether to run inference on this frame
            if frame_idx % sample_every_n == 0:
                # Save frame as an image temp file for the workflow
                with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp_img:
                    cv2.imwrite(tmp_img.name, frame)
                    result = _run_on_image_path(tmp_img.name, use_cache=use_cache)
                preds = _extract_predictions(result)
                last_preds = preds
                summary["workflow_calls"] += 1

                # Keep a small sample of results (avoid huge JSON)
                if len(summary["example_results"]) < 3:
                    summary["example_results"].append(result)

            # Draw using last_preds
            pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            pil = _draw_boxes_pil(pil, last_preds, conf_thresh=conf_thresh)
            out_frame = cv2.cvtColor(np.array(pil), cv2.COLOR_RGB2BGR)

            writer.write(out_frame)
            summary["frames_processed"] += 1
            frame_idx += 1

    except HTTPCallErrorError as e:
        return None, {
            "error": "Roboflow request failed during video processing",
            "status_code": getattr(e, "status_code", None),
            "api_message": getattr(e, "api_message", str(e)),
            "description": str(e),
        }
    finally:
        cap.release()
        writer.release()

    return tmp_out.name, summary


# ---- Gradio UI ----
with gr.Blocks(title="Roboflow Workflow Runner (Image + Video)") as demo:
    gr.Markdown("# Roboflow Workflow Runner (Image + Video)\nUpload an image or a video, run your workflow, and see bounding boxes.")

    with gr.Tab("Image"):
        img_in = gr.Image(type="filepath", label="Upload an image")
        img_cache = gr.Checkbox(value=True, label="Use cache (faster for repeat requests)")
        img_conf = gr.Slider(0.0, 1.0, value=0.25, step=0.05, label="Confidence threshold")

        img_btn = gr.Button("Run on Image")
        img_out = gr.Image(type="pil", label="Annotated image")
        img_json = gr.JSON(label="Raw workflow result")

        img_btn.click(fn=infer_image, inputs=[img_in, img_cache, img_conf], outputs=[img_out, img_json])

    with gr.Tab("Video"):
        vid_in = gr.Video(label="Upload a video")
        vid_cache = gr.Checkbox(value=True, label="Use cache (usually OFF for video, but you can try)")
        vid_conf = gr.Slider(0.0, 1.0, value=0.25, step=0.05, label="Confidence threshold")
        sample_every_n = gr.Slider(1, 30, value=5, step=1, label="Run inference every N frames (higher = cheaper/faster)")
        fps_out = gr.Slider(5, 30, value=15, step=1, label="Output video FPS")

        vid_btn = gr.Button("Run on Video")
        vid_out = gr.Video(label="Annotated video")
        vid_summary = gr.JSON(label="Video summary (includes a few sample results)")

        vid_btn.click(
            fn=infer_video,
            inputs=[vid_in, vid_cache, vid_conf, fps_out, sample_every_n],
            outputs=[vid_out, vid_summary]
        )

demo.launch()