app.py

import base64
import html
import json
import os
import shutil
import subprocess
import sys
import threading
from collections import Counter
from pathlib import Path

import spaces
import gradio as gr
import numpy as np
from PIL import Image, ImageDraw

ROOT = Path(__file__).resolve().parent
# Hugging Face Spaces local disk is ephemeral. If a persistent volume is
# attached, HF mounts it at /data; use it automatically so saved images,
# annotations, and models survive Space restarts/rebuilds.
PERSISTENT_ROOT = Path("/data")
if PERSISTENT_ROOT.exists() and os.access(PERSISTENT_ROOT, os.W_OK):
    DATA = PERSISTENT_ROOT / "anything_detector"
else:
    DATA = ROOT / "data"
IMAGE_DIR = DATA / "images"
DB_FILE = DATA / "dataset.json"

MODEL_DIR = ROOT / "model"
CLASSES_FILE = ROOT / "training" / "classes.txt"

IMAGE_DIR.mkdir(parents=True, exist_ok=True)
MODEL_DIR.mkdir(parents=True, exist_ok=True)

training_state = {
    "running": False,
    "message": "not started",
    "error": None,
}

detector_model = None
detector_processor = None


def get_classes():
    if not CLASSES_FILE.exists():
        return []
    return [
        x.strip()
        for x in CLASSES_FILE.read_text(encoding="utf-8").splitlines()
        if x.strip()
    ]


def load_db():
    if not DB_FILE.exists():
        return {"images": []}
    return json.loads(DB_FILE.read_text(encoding="utf-8"))


def save_db(db):
    DATA.mkdir(parents=True, exist_ok=True)
    tmp = DB_FILE.with_suffix(".json.tmp")
    tmp.write_text(
        json.dumps(db, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )
    tmp.replace(DB_FILE)


def image_record(filename):
    return next(
        (x for x in load_db()["images"] if x["filename"] == filename),
        None,
    )


def gallery_items(db):
    """Return gallery items as PIL images, not filesystem URLs.

    Using PIL objects is intentional: Gradio can render them directly, so
    the gallery does not depend on its static-file serving rules. The actual
    image is still saved to IMAGE_DIR and recorded in dataset.json.
    """
    items = []
    for item in db.get("images", []):
        path = Path(item.get("path", ""))
        if not path.is_file():
            continue
        try:
            with Image.open(path) as im:
                # Make a detached copy because the file handle closes after
                # leaving the context manager. Convert unusual modes so the
                # Gradio serializer can always handle the image.
                if im.mode not in ("RGB", "RGBA"):
                    im = im.convert("RGB")
                else:
                    im = im.copy()
                items.append((im, item.get("filename", path.name)))
        except Exception:
            # A bad file should not prevent the other saved images from
            # appearing in the gallery.
            continue
    return items


def normalize_uploaded_file(file):
    """Get a real local path from Gradio filepath/FileData variants."""
    if isinstance(file, (str, os.PathLike)):
        return Path(file)
    if isinstance(file, dict):
        value = file.get("path") or file.get("name") or file.get("url")
        if value and not str(value).startswith(("http://", "https://")):
            return Path(str(value))
    value = getattr(file, "path", None) or getattr(file, "name", None)
    if value:
        return Path(str(value))
    raise ValueError(f"Unsupported upload object: {type(file).__name__}")


def upload_training_images(files):
    """Persist uploaded images and refresh every dataset UI component."""
    DATA.mkdir(parents=True, exist_ok=True)
    IMAGE_DIR.mkdir(parents=True, exist_ok=True)
    db = load_db()
    if not isinstance(db, dict) or not isinstance(db.get("images"), list):
        db = {"images": []}
    if files is None:
        files = []
    elif not isinstance(files, list):
        files = [files]

    print(f"[SAVE IMAGES] received {len(files)} file(s)", flush=True)
    print(f"[SAVE IMAGES] DATA={DATA} IMAGE_DIR={IMAGE_DIR}", flush=True)
    saved, errors = [], []
    existing = {str(x.get("filename", "")) for x in db["images"]}

    for file in files:
        try:
            src = normalize_uploaded_file(file)
            print(f"[SAVE IMAGES] source={src}", flush=True)
            if not src.exists() or not src.is_file():
                raise FileNotFoundError(f"Uploaded file does not exist: {src}")
            ext = src.suffix.lower()
            if ext not in {".jpg", ".jpeg", ".png", ".webp", ".bmp"}:
                raise ValueError(f"Unsupported image type: {ext or 'unknown'}")
            with Image.open(src) as im:
                im.verify()
            with Image.open(src) as im:
                width, height = im.size

            dst = IMAGE_DIR / src.name
            stem = src.stem
            n = 1
            while dst.exists() or dst.name in existing:
                dst = IMAGE_DIR / f"{stem}_{n}{ext}"
                n += 1
            dst.write_bytes(src.read_bytes())
            if not dst.exists() or dst.stat().st_size == 0:
                raise IOError(f"Image copy failed: {dst}")
            db["images"].append({"filename": dst.name, "path": str(dst), "width": width, "height": height, "annotations": []})
            existing.add(dst.name)
            saved.append(dst.name)
            print(f"[SAVE IMAGES] SAVED {dst}", flush=True)
        except Exception as exc:
            errors.append(f"{getattr(file, 'name', file)}: {exc}")
            print(f"[SAVE IMAGES] ERROR: {exc}", flush=True)

    save_db(db)
    print(f"[SAVE IMAGES] database={DB_FILE} records={len(db['images'])}", flush=True)
    choices = [x["filename"] for x in db["images"]]
    latest = saved[-1] if saved else (choices[-1] if choices else None)
    location = "Persistent /data storage" if DATA == PERSISTENT_ROOT / "anything_detector" else "Space local storage"
    verified = sum(1 for x in db["images"] if Path(x.get("path", "")).is_file())
    if saved and not errors:
        msg = f"### ✅ Saved {len(saved)} image(s)\n\n**Dataset:** `{DB_FILE}`  \n**Image files:** `{IMAGE_DIR}`  \n**Verified on disk:** {verified}  \n**Storage:** {location}"
    elif saved:
        msg = f"### ⚠️ Saved {len(saved)} image(s), with {len(errors)} error(s)\n\n**Verified on disk:** {verified}\n\n" + "\n".join(f"- {e}" for e in errors)
    elif errors:
        msg = "### ❌ No images were saved\n\n" + "\n".join(f"- {e}" for e in errors)
    else:
        msg = "### ⚠️ No images selected."
    gallery_value = gallery_items(db)
    annotator_value = annotation_html(latest) if latest else annotation_html(None)
    return (gr.update(choices=choices, value=latest), msg, gallery_value, gr.update(choices=choices, value=latest), annotator_value, gallery_value)

def make_data_url(path):
    mime = "image/png"
    suffix = Path(path).suffix.lower()
    if suffix in [".jpg", ".jpeg"]:
        mime = "image/jpeg"
    elif suffix == ".webp":
        mime = "image/webp"

    data = base64.b64encode(Path(path).read_bytes()).decode("ascii")
    return f"data:{mime};base64,{data}"


def annotation_html(filename):
    if not filename:
        return """
        <div style="padding:30px;border:1px dashed #999;border-radius:12px">
            Select a training image to start annotating.
        </div>
        """

    item = image_record(filename)
    if not item:
        return "<div>Image not found.</div>"

    image_url = make_data_url(item["path"])
    classes_json = json.dumps(get_classes())
    existing = []

    for ann in item.get("annotations", []):
        x, y, w, h = ann["box"]
        existing.append({
            "x": x,
            "y": y,
            "w": w,
            "h": h,
            "label": ann["class"],
        })

    existing_json = json.dumps(existing)

    # The canvas is scaled visually, but boxes are stored in original image pixels.
    # This makes annotations independent of browser size.
    return f"""
    <div class="annotator-root"
         data-image="{html.escape(image_url, quote=True)}"
         data-classes='{html.escape(classes_json, quote=True)}'
         data-existing='{html.escape(existing_json, quote=True)}'
         style="font-family:system-ui; width:100%;">
      <div style="display:flex;gap:8px;align-items:center;flex-wrap:wrap;margin-bottom:10px">
        <select class="ann-class" style="padding:8px;border-radius:8px;border:1px solid #aaa"></select>
        <button type="button" class="ann-mode" style="padding:8px 12px;border-radius:8px;border:1px solid #aaa;background:#eee">Draw boxes</button>
        <button type="button" class="ann-delete" style="padding:8px 12px;border-radius:8px;border:1px solid #aaa">Delete selected</button>
        <button type="button" class="ann-clear" style="padding:8px 12px;border-radius:8px;border:1px solid #aaa">Clear all</button>
        <button type="button" class="ann-save" style="padding:8px 14px;border-radius:8px;border:0;background:#2563eb;color:white;font-weight:600">Save annotations</button>
        <span class="ann-info" style="font-size:13px;color:#555">Drag on the image to draw a box.</span>
      </div>
      <div style="position:relative;width:100%;overflow:auto;border:1px solid #ccc;border-radius:10px;background:#111;">
        <canvas class="ann-canvas" style="display:block;max-width:100%;height:auto;cursor:crosshair;"></canvas>
      </div>
      <div class="ann-count" style="margin-top:8px;font-size:13px;color:#555"></div>
    </div>
    """


ANNOTATOR_JS = r"""
function setupAnythingAnnotator() {
  const root = document.querySelector(".annotator-root:not([data-ann-initialized])");
  if (!root) return;
  root.dataset.annInitialized = "1";

  const canvas = root.querySelector(".ann-canvas");
  const ctx = canvas.getContext("2d");
  const img = new Image();

  let boxes = JSON.parse(root.dataset.existing || "[]");
  let selected = -1;
  let drawing = false;
  let start = null;
  let draft = null;

  const classes = JSON.parse(root.dataset.classes || "[]");
  const select = root.querySelector(".ann-class");
  const count = root.querySelector(".ann-count");
  const info = root.querySelector(".ann-info");

  select.innerHTML = "";
  classes.forEach(c => {
    const o = document.createElement("option");
    o.value = c;
    o.textContent = c;
    select.appendChild(o);
  });

  function resizeCanvas() {
    canvas.width = img.naturalWidth;
    canvas.height = img.naturalHeight;
    render();
  }

  function pointerPosition(e) {
    const r = canvas.getBoundingClientRect();
    const sx = canvas.width / r.width;
    const sy = canvas.height / r.height;
    return {
      x: Math.max(0, Math.min(canvas.width, (e.clientX - r.left) * sx)),
      y: Math.max(0, Math.min(canvas.height, (e.clientY - r.top) * sy))
    };
  }

  function render() {
    if (!img.naturalWidth) return;

    ctx.clearRect(0, 0, canvas.width, canvas.height);
    ctx.drawImage(img, 0, 0, canvas.width, canvas.height);

    boxes.forEach((b, i) => {
      ctx.lineWidth = i === selected ? 5 : 3;
      ctx.strokeStyle = i === selected ? "#00e5ff" : "#ff3030";
      ctx.fillStyle = i === selected
        ? "rgba(0,229,255,.12)"
        : "rgba(255,48,48,.10)";
      ctx.fillRect(b.x, b.y, b.w, b.h);
      ctx.strokeRect(b.x, b.y, b.w, b.h);

      ctx.font = "bold 16px sans-serif";
      const text = `${b.label}`;
      const tw = ctx.measureText(text).width + 10;
      ctx.fillStyle = i === selected ? "#00a8c7" : "#ff3030";
      ctx.fillRect(b.x, Math.max(0, b.y - 24), tw, 24);
      ctx.fillStyle = "white";
      ctx.fillText(text, b.x + 5, Math.max(17, b.y - 7));

      if (i === selected) {
        ctx.fillStyle = "#00e5ff";
        [[b.x,b.y],[b.x+b.w,b.y],[b.x,b.y+b.h],[b.x+b.w,b.y+b.h]].forEach(p => {
          ctx.fillRect(p[0]-5,p[1]-5,10,10);
        });
      }
    });

    if (draft) {
      ctx.setLineDash([8,5]);
      ctx.lineWidth = 3;
      ctx.strokeStyle = "#00ff66";
      ctx.strokeRect(draft.x, draft.y, draft.w, draft.h);
      ctx.setLineDash([]);
    }

    count.textContent = `${boxes.length} box(es)`;
  }

  function hitTest(p) {
    for (let i = boxes.length - 1; i >= 0; i--) {
      const b = boxes[i];
      if (p.x >= b.x && p.x <= b.x+b.w &&
          p.y >= b.y && p.y <= b.y+b.h) return i;
    }
    return -1;
  }

  canvas.onpointerdown = e => {
    const p = pointerPosition(e);
    selected = hitTest(p);

    // If clicking an existing box, select it; otherwise begin a new box.
    if (selected >= 0) {
      select.value = boxes[selected].label;
      render();
      return;
    }

    drawing = true;
    start = p;
    draft = {x:p.x,y:p.y,w:0,h:0,label:select.value};
    canvas.setPointerCapture(e.pointerId);
    render();
  };

  canvas.onpointermove = e => {
    if (!drawing) return;
    const p = pointerPosition(e);
    draft = {
      x: Math.min(start.x,p.x),
      y: Math.min(start.y,p.y),
      w: Math.abs(p.x-start.x),
      h: Math.abs(p.y-start.y),
      label: select.value
    };
    render();
  };

  canvas.onpointerup = e => {
    if (!drawing) return;
    drawing = false;
    canvas.releasePointerCapture(e.pointerId);

    if (draft && draft.w >= 5 && draft.h >= 5) {
      boxes.push(draft);
      selected = boxes.length - 1;
    }

    draft = null;
    render();
  };

  select.onchange = () => {
    if (selected >= 0) {
      boxes[selected].label = select.value;
      render();
    }
  };

  root.querySelector(".ann-delete").onclick = () => {
    if (selected >= 0) {
      boxes.splice(selected,1);
      selected = -1;
      render();
    }
  };

  root.querySelector(".ann-clear").onclick = () => {
    boxes = [];
    selected = -1;
    render();
  };

  root.querySelector(".ann-mode").onclick = () => {
    selected = -1;
    info.textContent = "Draw mode: drag around each object.";
    canvas.style.cursor = "crosshair";
    render();
  };

  root.querySelector(".ann-save").onclick = () => {
    // Send annotation JSON through the hidden textbox.
    const payload = JSON.stringify({
      filename: root.closest(".annotator-shell").dataset.filename,
      boxes: boxes
    });

    const input = document.querySelector("#annotation_payload textarea, #annotation_payload input");
    if (!input) {
      alert("Annotation output field was not found. Reload the Space.");
      return;
    }

    const setter = Object.getOwnPropertyDescriptor(
      HTMLTextAreaElement.prototype, "value"
    )?.set || Object.getOwnPropertyDescriptor(
      HTMLInputElement.prototype, "value"
    )?.set;

    if (setter) setter.call(input, payload);
    else input.value = payload;

    input.dispatchEvent(new Event("input", {bubbles:true}));
    input.dispatchEvent(new Event("change", {bubbles:true}));
  };

  img.onload = resizeCanvas;
  img.src = root.dataset.image;

  window.addEventListener("resize", render);
}

setupAnythingAnnotator();

// Gradio replaces the HTML component's DOM when an image is selected.
// Re-initialize the canvas whenever a new annotator root appears.
const anythingAnnotatorObserver = new MutationObserver(() => {
  setupAnythingAnnotator();
});
anythingAnnotatorObserver.observe(document.body, { childList: true, subtree: true });
"""


def save_annotation_payload(payload):
    if not payload:
        return "No annotation data received."

    try:
        data = json.loads(payload)
        filename = data["filename"]
        boxes = data.get("boxes", [])
    except Exception as e:
        return f"Invalid annotation data: {e}"

    item = image_record(filename)
    if not item:
        return "Image not found."

    valid = set(get_classes())
    saved = []

    for b in boxes:
        label = str(b.get("label", "")).strip()
        if label not in valid:
            continue

        x = max(0, float(b.get("x", 0)))
        y = max(0, float(b.get("y", 0)))
        w = max(0, float(b.get("w", 0)))
        h = max(0, float(b.get("h", 0)))

        if w < 2 or h < 2:
            continue

        x = min(x, item["width"] - 1)
        y = min(y, item["height"] - 1)
        w = min(w, item["width"] - x)
        h = min(h, item["height"] - y)

        saved.append({
            "class": label,
            "box": [x, y, w, h],
        })

    db = load_db()
    target = next(
        x for x in db["images"]
        if x["filename"] == filename
    )
    target["annotations"] = saved
    save_db(db)

    return f"Saved {len(saved)} box(es) for **{filename}**."


def build_coco_dataset():
    db = load_db()
    classes = get_classes()

    items = [
        x for x in db["images"]
        if x.get("annotations")
    ]

    if len(items) < 2:
        raise RuntimeError("Annotate at least 2 images before training.")

    generated = DATA / "generated"
    if generated.exists():
        shutil.rmtree(generated)

    train_dir = generated / "train"
    val_dir = generated / "val"
    train_dir.joinpath("images").mkdir(parents=True)
    val_dir.joinpath("images").mkdir(parents=True)

    split = max(1, int(len(items) * 0.8))
    if split >= len(items):
        split = len(items) - 1

    cat = {name: i + 1 for i, name in enumerate(classes)}

    def make_split(selected, outdir):
        images = []
        annotations = []
        ann_id = 1

        for item in selected:
            shutil.copy2(
                item["path"],
                outdir / "images" / item["filename"]
            )

            images.append({
                "id": item["filename"],
                "file_name": item["filename"],
                "width": item["width"],
                "height": item["height"],
            })

            for a in item["annotations"]:
                x, y, w, h = a["box"]
                annotations.append({
                    "id": ann_id,
                    "image_id": item["filename"],
                    "category_id": cat[a["class"]],
                    "bbox": [x, y, w, h],
                    "area": w * h,
                    "iscrowd": 0,
                })
                ann_id += 1

        coco = {
            "images": images,
            "annotations": annotations,
            "categories": [
                {"id": i+1, "name": name}
                for i, name in enumerate(classes)
            ],
        }

        (outdir / "annotations.json").write_text(
            json.dumps(coco, indent=2),
            encoding="utf-8",
        )

    make_split(items[:split], train_dir)
    make_split(items[split:], val_dir)

    return train_dir, val_dir


def train_worker():
    global training_state, detector_model, detector_processor

    try:
        training_state = {
            "running": True,
            "message": "Building COCO dataset...",
            "error": None,
        }

        train_dir, val_dir = build_coco_dataset()

        training_state["message"] = "Training RT-DETR..."

        command = [
            sys.executable,
            str(ROOT / "training" / "train.py"),
            "--train-dir", str(train_dir),
            "--val-dir", str(val_dir),
            "--classes", str(CLASSES_FILE),
            "--output-dir", str(MODEL_DIR),
            "--epochs", os.getenv("EPOCHS", "30"),
            "--batch-size", os.getenv("BATCH_SIZE", "2"),
        ]

        result = subprocess.run(
            command,
            cwd=ROOT,
            capture_output=True,
            text=True,
        )

        if result.returncode != 0:
            raise RuntimeError(
                (result.stderr or result.stdout)[-12000:]
            )

        detector_model = None
        detector_processor = None

        training_state = {
            "running": False,
            "message": "Training complete.",
            "error": None,
        }

    except Exception as exc:
        training_state = {
            "running": False,
            "message": "Training failed.",
            "error": str(exc),
        }


@spaces.GPU(duration=1500)
def start_training():
    """Train the RT-DETR detector using the Space GPU."""
    if training_state["running"]:
        return "Training is already running."

    # Run synchronously while the ZeroGPU allocation is active.
    # Starting a background thread here would release the GPU as soon as
    # this function returned, leaving the training subprocess without it.
    train_worker()
    return get_training_status()


def get_training_status():
    return json.dumps(training_state, indent=2)


def load_detector():
    global detector_model, detector_processor

    if not (MODEL_DIR / "config.json").exists():
        raise RuntimeError("No trained model found. Train first.")

    if detector_model is None:
        from transformers import (
            RTDetrForObjectDetection,
            RTDetrImageProcessor,
        )

        detector_processor = (
            RTDetrImageProcessor.from_pretrained(str(MODEL_DIR))
        )
        detector_model = (
            RTDetrForObjectDetection.from_pretrained(str(MODEL_DIR))
        )

        import torch
        device = "cuda" if torch.cuda.is_available() else "cpu"
        detector_model.to(device)
        detector_model.eval()

    return detector_processor, detector_model


@spaces.GPU(duration=300)
def detect(image, threshold):
    if image is None:
        return None, "Upload an image."

    try:
        processor, model = load_detector()
    except Exception as e:
        return None, str(e)

    import torch

    pil = (
        Image.fromarray(image.astype("uint8")).convert("RGB")
        if isinstance(image, np.ndarray)
        else image.convert("RGB")
    )

    device = next(model.parameters()).device
    inputs = processor(images=pil, return_tensors="pt")
    inputs = {
        k: v.to(device) if torch.is_tensor(v) else v
        for k, v in inputs.items()
    }

    with torch.inference_mode():
        outputs = model(**inputs)

    target_sizes = torch.tensor(
        [[pil.height, pil.width]],
        device=device,
    )

    result = processor.post_process_object_detection(
        outputs,
        threshold=float(threshold),
        target_sizes=target_sizes,
    )[0]

    detections = []
    counts = Counter()

    for score, label, box in zip(
        result["scores"],
        result["labels"],
        result["boxes"],
    ):
        cls = model.config.id2label[int(label.item())]
        conf = float(score.item())
        coords = [round(float(x), 2) for x in box.tolist()]

        detections.append({
            "class": cls,
            "confidence": round(conf, 4),
            "box": coords,
        })
        counts[cls] += 1

    out = pil.copy()
    draw = ImageDraw.Draw(out)

    for d in detections:
        x1, y1, x2, y2 = d["box"]
        draw.rectangle(
            [x1, y1, x2, y2],
            outline=(255, 0, 0),
            width=4,
        )
        text = f'{d["class"]} {d["confidence"]:.2f}'
        draw.text((x1 + 4, max(2, y1 + 4)), text, fill=(255, 0, 0))

    response = {
        "total": len(detections),
        "counts": dict(sorted(counts.items())),
        "detections": detections,
    }

    return out, json.dumps(response, indent=2)


with gr.Blocks(title="Anything Detector", js=ANNOTATOR_JS) as demo:
    gr.Markdown(
        "# 🔍 Anything Detector\n"
        "Train your own object detector for any classes you define. **No YOLO.**"
    )

    with gr.Tab("1 · Upload training images"):
        files = gr.File(
            file_count="multiple",
            file_types=["image"],
            type="filepath",
            label="Training images",
        )
        save_images = gr.Button(
            "Save training images",
            variant="primary",
        )
        image_selector = gr.Dropdown(
            choices=[
                x["filename"]
                for x in load_db()["images"]
            ],
            label="Training image",
        )
        upload_status = gr.Markdown()
        saved_images = gr.Gallery(
            value=gallery_items(load_db()),
            label="Saved training images",
            columns=5,
            rows=3,
            height=420,
            object_fit="contain",
            show_label=True,
            allow_preview=True,
        )


    with gr.Tab("2 · Draw boxes"):
        gr.Markdown(
            """
### Draw boxes with your mouse

1. Select an image.
2. The full image appears.
3. Select the object class.
4. **Click and drag around each object.**
5. Click a box to select it.
6. Change its class if needed.
7. Click **Save annotations**.

This is a native HTML5 canvas — there is no `gradio_image_annotation` package.
"""
        )

        _initial_images = load_db()["images"]
        _initial_filename = _initial_images[0]["filename"] if _initial_images else None
        ann_selector = gr.Dropdown(
            choices=[x["filename"] for x in _initial_images],
            value=_initial_filename,
            label="Image to annotate",
        )

        # The shell is deliberately separate from gr.HTML so JS can identify
        # the current filename without depending on Gradio's internal DOM.
        annotator_html = gr.HTML(
            value=annotation_html(
                (
                    load_db()["images"][0]["filename"]
                    if load_db()["images"]
                    else None
                )
            ),
            elem_id="annotator_html",
        )

        # Hidden data channel used by the canvas Save button.
        annotation_payload = gr.Textbox(
            value="",
            visible=False,
            elem_id="annotation_payload",
        )

        ann_save_status = gr.Markdown()

        def update_annotator(filename):
            content = annotation_html(filename)
            # Add a filename marker to the root so the JS knows what it is.
            marker = (
                f'<div class="annotator-shell" '
                f'data-filename="{html.escape(filename or "", quote=True)}">'
            )
            return gr.HTML(
                value=marker + content + "</div>"
            ).value

        ann_selector.change(
            lambda filename: (
                f'<div class="annotator-shell" '
                f'data-filename="{html.escape(filename or "", quote=True)}">'
                f'{annotation_html(filename)}'
                f'</div>'
            ),
            inputs=ann_selector,
            outputs=annotator_html,
        )

        save_images.click(
            upload_training_images,
            inputs=files,
            outputs=[image_selector, upload_status, saved_images, ann_selector, annotator_html],
        )

        annotation_payload.change(
            save_annotation_payload,
            inputs=annotation_payload,
            outputs=ann_save_status,
        )

    with gr.Tab("3 · Dataset"):
        stats = gr.Markdown()
        dataset_images = gr.Gallery(value=gallery_items(load_db()), label="Images currently in dataset", columns=5, rows=3, height=420, object_fit="contain")
        refresh = gr.Button("Refresh")
        refresh.click(
            lambda: (
                f"Images: {len(load_db()['images'])}  \n"
                f"Annotated images: {sum(bool(x.get('annotations')) for x in load_db()['images'])}  \n"
                f"Boxes: {sum(len(x.get('annotations', [])) for x in load_db()['images'])}  \n"
                f"Storage: {'Persistent /data' if DATA == PERSISTENT_ROOT / 'anything_detector' else 'Temporary Space storage'}  \n"
                f"Classes: {', '.join(get_classes())}",
                gallery_items(load_db()),
            ),
            outputs=[stats, dataset_images],
        )

    with gr.Tab("4 · Train RT-DETR"):
        gr.Markdown(
            "Train from the annotations you saved. Use GPU hardware in Hugging Face for serious datasets."
        )
        train_button = gr.Button(
            "🚀 Start training",
            variant="primary",
        )
        status_button = gr.Button("Refresh status")
        status = gr.Code(language="json")

        train_button.click(
            start_training,
            outputs=status,
        )
        status_button.click(
            get_training_status,
            outputs=status,
        )

    with gr.Tab("5 · Detect"):
        input_image = gr.Image(
            type="numpy",
            label="New image",
        )
        threshold = gr.Slider(
            minimum=0.05,
            maximum=0.95,
            value=0.35,
            step=0.05,
            label="Confidence threshold",
        )
        detect_button = gr.Button(
            "🔍 Detect and count",
            variant="primary",
        )
        result_image = gr.Image(
            type="pil",
            label="Detection result",
        )
        result = gr.Code(
            language="json",
            label="Counts and detections",
        )

        detect_button.click(
            detect,
            inputs=[input_image, threshold],
            outputs=[result_image, result],
        )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        ssr_mode=False,
        max_file_size="100mb",
        allowed_paths=[str(DATA), str(IMAGE_DIR)],
    )