app.py

from __future__ import annotations

import os
from functools import lru_cache
from pathlib import Path

import gradio as gr
from PIL import Image, ImageDraw, ImageFont
from ultralytics import YOLO


ROOT = Path(__file__).resolve().parent
MODEL_PATH = ROOT / "models" / "best.pt"
MODEL_URL = "https://huggingface.co/DefendIntelligence/vessel-detection/resolve/main/models/best.pt"
EXAMPLES_DIR = ROOT / "examples"
MAIN_EXAMPLE_PATH = EXAMPLES_DIR / "example-00-multi-vessel-patch.png"
MAX_TILES = 196
BATCH_SIZE = 8


@lru_cache(maxsize=1)
def load_model() -> YOLO:
    if not MODEL_PATH.exists():
        raise FileNotFoundError(
            f"Model not found: {MODEL_PATH}. Run `python run_local.py` or download it from {MODEL_URL}."
        )
    return YOLO(str(MODEL_PATH))


def _tile_starts(length: int, tile_size: int, overlap: int) -> list[int]:
    if length <= tile_size:
        return [0]
    stride = max(1, tile_size - overlap)
    starts = list(range(0, max(1, length - tile_size + 1), stride))
    last = length - tile_size
    if starts[-1] != last:
        starts.append(last)
    return starts


def _iter_tiles(image: Image.Image, tile_size: int, overlap: int) -> list[tuple[Image.Image, int, int]]:
    width, height = image.size
    x_starts = _tile_starts(width, tile_size, overlap)
    y_starts = _tile_starts(height, tile_size, overlap)
    tiles: list[tuple[Image.Image, int, int]] = []
    for y in y_starts:
        for x in x_starts:
            right = min(width, x + tile_size)
            bottom = min(height, y + tile_size)
            tiles.append((image.crop((x, y, right, bottom)), x, y))
    return tiles


def _box_iou(a: list[float], b: list[float]) -> float:
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = b
    inter_x1 = max(ax1, bx1)
    inter_y1 = max(ay1, by1)
    inter_x2 = min(ax2, bx2)
    inter_y2 = min(ay2, by2)
    inter_w = max(0.0, inter_x2 - inter_x1)
    inter_h = max(0.0, inter_y2 - inter_y1)
    inter_area = inter_w * inter_h
    if inter_area <= 0:
        return 0.0
    area_a = max(0.0, ax2 - ax1) * max(0.0, ay2 - ay1)
    area_b = max(0.0, bx2 - bx1) * max(0.0, by2 - by1)
    union = area_a + area_b - inter_area
    return inter_area / union if union > 0 else 0.0


def _nms(detections: list[dict], iou_threshold: float) -> list[dict]:
    remaining = sorted(detections, key=lambda item: float(item["confidence"]), reverse=True)
    kept: list[dict] = []
    while remaining:
        current = remaining.pop(0)
        kept.append(current)
        remaining = [
            item
            for item in remaining
            if item["class_id"] != current["class_id"]
            or _box_iou(item["box"], current["box"]) < iou_threshold
        ]
    return kept


def _model_names(model: YOLO) -> dict[int, str]:
    names = getattr(model, "names", None) or {}
    if isinstance(names, dict):
        return {int(key): str(value) for key, value in names.items()}
    return {index: str(name) for index, name in enumerate(names)}


def _predict_tiles(
    image: Image.Image,
    *,
    confidence: float,
    iou: float,
    tile_size: int,
    overlap: int,
    max_det: int,
) -> tuple[list[dict], int]:
    model = load_model()
    names = _model_names(model)
    rgb_image = image.convert("RGB")
    safe_tile_size = max(320, int(tile_size))
    safe_overlap = max(0, min(int(overlap), safe_tile_size - 32))
    tiles = _iter_tiles(rgb_image, safe_tile_size, safe_overlap)

    if len(tiles) > MAX_TILES:
        raise ValueError(
            f"Image too large for this CPU Space: {len(tiles)} tiles. "
            f"Resize the image or increase the tile size."
        )

    detections: list[dict] = []
    for start in range(0, len(tiles), BATCH_SIZE):
        batch = tiles[start : start + BATCH_SIZE]
        batch_images = [tile for tile, _, _ in batch]
        results = model.predict(
            source=batch_images,
            conf=float(confidence),
            iou=float(iou),
            imgsz=safe_tile_size,
            max_det=int(max_det),
            verbose=False,
        )
        for result, (_, offset_x, offset_y) in zip(results, batch):
            boxes = getattr(result, "boxes", None)
            if boxes is None or len(boxes) == 0:
                continue
            xyxy = boxes.xyxy.cpu().numpy()
            confs = boxes.conf.cpu().numpy()
            classes = boxes.cls.cpu().numpy().astype(int)
            for box, score, class_id in zip(xyxy, confs, classes):
                x1, y1, x2, y2 = box.tolist()
                detections.append(
                    {
                        "label": names.get(int(class_id), f"class_{int(class_id)}"),
                        "class_id": int(class_id),
                        "confidence": float(score),
                        "box": [
                            float(x1 + offset_x),
                            float(y1 + offset_y),
                            float(x2 + offset_x),
                            float(y2 + offset_y),
                        ],
                    }
                )

    detections = _nms(detections, float(iou))
    detections = detections[: int(max_det)]
    return detections, len(tiles)


def _draw_detections(image: Image.Image, detections: list[dict]) -> Image.Image:
    annotated = image.convert("RGB").copy()
    draw = ImageDraw.Draw(annotated)
    font = ImageFont.load_default()
    line_width = max(2, round(max(annotated.size) / 420))

    for detection in detections:
        x1, y1, x2, y2 = detection["box"]
        label = f"{detection['label']} {detection['confidence']:.2f}"
        draw.rectangle((x1, y1, x2, y2), outline=(255, 64, 48), width=line_width)
        text_box = draw.textbbox((x1, y1), label, font=font)
        text_w = text_box[2] - text_box[0]
        text_h = text_box[3] - text_box[1]
        label_y = max(0, y1 - text_h - 6)
        draw.rectangle((x1, label_y, x1 + text_w + 8, label_y + text_h + 6), fill=(255, 64, 48))
        draw.text((x1 + 4, label_y + 3), label, fill=(255, 255, 255), font=font)

    return annotated


def _table_rows(detections: list[dict]) -> list[list[object]]:
    rows: list[list[object]] = []
    for index, detection in enumerate(detections, start=1):
        x1, y1, x2, y2 = detection["box"]
        rows.append(
            [
                index,
                detection["label"],
                round(float(detection["confidence"]), 4),
                round(x1, 1),
                round(y1, 1),
                round(x2, 1),
                round(y2, 1),
                round(x2 - x1, 1),
                round(y2 - y1, 1),
            ]
        )
    return rows


def detect_boats(
    image: Image.Image | None,
    confidence: float,
    iou: float,
    tile_size: int,
    overlap: int,
    max_det: int,
) -> tuple[Image.Image | None, list[list[object]], str]:
    if image is None:
        return None, [], "Upload a satellite image to run detection."

    try:
        detections, tile_count = _predict_tiles(
            image,
            confidence=confidence,
            iou=iou,
            tile_size=tile_size,
            overlap=overlap,
            max_det=max_det,
        )
    except Exception as exc:
        return image, [], f"Inference error: {exc}"

    annotated = _draw_detections(image, detections)
    rows = _table_rows(detections)
    if detections:
        summary = f"{len(detections)} detection(s) above {confidence:.2f}. Tiles analyzed: {tile_count}."
    else:
        summary = f"No detections above {confidence:.2f}. Tiles analyzed: {tile_count}."
    return annotated, rows, summary


def _example_paths() -> list[list[str]]:
    paths = sorted(EXAMPLES_DIR.glob("*.png"))
    return [[str(path)] for path in paths[:10]]


with gr.Blocks(title="Vessel Detection") as demo:
    gr.Markdown(
        """
        # Vessel Detection

        Fine-tuned YOLOv8s model for detecting vessels in RGB satellite imagery.
        Upload a satellite image or select an example, then run detection.
        """
    )
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(
                value=str(MAIN_EXAMPLE_PATH) if MAIN_EXAMPLE_PATH.exists() else None,
                type="pil",
                label="Satellite image",
            )
            confidence_input = gr.Slider(0.01, 0.95, value=0.20, step=0.01, label="Confidence threshold")
            iou_input = gr.Slider(0.05, 0.90, value=0.45, step=0.05, label="IoU NMS")
            tile_size_input = gr.Slider(320, 1024, value=640, step=32, label="Tile size")
            overlap_input = gr.Slider(0, 256, value=96, step=16, label="Tile overlap")
            max_det_input = gr.Slider(1, 200, value=80, step=1, label="Max detections")
            run_button = gr.Button("Detect vessels", variant="primary")
        with gr.Column(scale=1):
            output_image = gr.Image(type="pil", label="Annotated image")
            summary_output = gr.Markdown()
            table_output = gr.Dataframe(
                headers=["#", "label", "confidence", "x1", "y1", "x2", "y2", "width", "height"],
                datatype=["number", "str", "number", "number", "number", "number", "number", "number", "number"],
                label="Detections",
            )

    run_button.click(
        fn=detect_boats,
        inputs=[image_input, confidence_input, iou_input, tile_size_input, overlap_input, max_det_input],
        outputs=[output_image, table_output, summary_output],
    )

    gr.Examples(
        examples=_example_paths(),
        inputs=[image_input],
        label="Example images",
    )


if __name__ == "__main__":
    launch_kwargs = {}
    if os.environ.get("GRADIO_SERVER_NAME"):
        launch_kwargs["server_name"] = os.environ["GRADIO_SERVER_NAME"]
    if os.environ.get("GRADIO_SERVER_PORT"):
        launch_kwargs["server_port"] = int(os.environ["GRADIO_SERVER_PORT"])
    demo.launch(**launch_kwargs)