app.py

"""
Gradio app: upload a Roboflow COCO-JSON train/test split, train a FOMO-style
object detection model, quantize it to int8 TFLite, test it, and download
the result for deployment on a Seeed XIAO ESP32S3 Sense.

Run locally with:  python app.py
On Hugging Face Spaces with ZeroGPU hardware selected, the training function
below is automatically given a real GPU for the duration of each run.
"""

import os
import math
import shutil
import tempfile
import json

import gradio as gr
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

import fomo_utils as fu

# --------------------------------------------------------------------------
# ZeroGPU support (Hugging Face Spaces). The `spaces` package is only
# installed/meaningful on HF Spaces; locally it's simply not present, so we
# fall back to a no-op decorator that behaves identically to how @spaces.GPU
# behaves outside a ZeroGPU environment anyway.
# --------------------------------------------------------------------------
try:
    import spaces
    HAS_ZEROGPU = True
except ImportError:
    HAS_ZEROGPU = False

    class _NoOpSpaces:
        @staticmethod
        def GPU(*deco_args, **deco_kwargs):
            # Support both @spaces.GPU and @spaces.GPU(duration=...) usage.
            if len(deco_args) == 1 and callable(deco_args[0]):
                return deco_args[0]

            def wrap(fn):
                return fn
            return wrap

    spaces = _NoOpSpaces()


WORK_DIR = os.path.join(tempfile.gettempdir(), "fomo_trainer")
os.makedirs(WORK_DIR, exist_ok=True)

# Fixed, known paths on disk rather than an in-memory session dict. ZeroGPU
# (and, in general, any setup where the training call might run in a
# separate worker process) doesn't guarantee that mutations to a Python
# global made inside the GPU-decorated function are visible back in the
# main process — but files written to the shared filesystem are. The Test
# and Download tabs read these paths directly instead of relying on shared
# in-memory state.
TFLITE_PATH = os.path.join(WORK_DIR, "crow_fomo_int8.tflite")
H_PATH = os.path.join(WORK_DIR, "crow_fomo_model.h")
LABELS_PATH = os.path.join(WORK_DIR, "labels.json")
PLOT_PATH = os.path.join(WORK_DIR, "loss_curve.png")


def _training_duration(train_files, test_files, input_size, epochs, batch_size,
                        learning_rate, background_weight, optimizer_name):
    """Rough estimate of GPU seconds needed, for ZeroGPU's queue/duration system.
    A small FOMO model trains fast on a real GPU; this pads generously for
    backbone download + dataset parsing + int8 quantization overhead."""
    per_epoch_seconds = 2.0
    overhead_seconds = 45
    estimate = int(int(epochs) * per_epoch_seconds + overhead_seconds)
    return max(30, min(estimate, 300))  # ZeroGPU allows requesting up to a few minutes


@spaces.GPU(duration=_training_duration)
def run_training(train_files, test_files, input_size, epochs, batch_size,
                  learning_rate, background_weight, optimizer_name):
    if not train_files or not test_files:
        yield "Please select both a train folder and a test folder first.", None, None
        return

    shutil.rmtree(WORK_DIR, ignore_errors=True)
    os.makedirs(WORK_DIR, exist_ok=True)

    log_lines = ["Parsing uploaded folders..."]
    yield "\n".join(log_lines), None, None

    try:
        train_samples, train_cats, train_missing = fu.load_coco_split_from_files(train_files, "train")
        test_samples, test_cats, test_missing = fu.load_coco_split_from_files(test_files, "test")
    except FileNotFoundError as e:
        yield str(e), None, None
        return

    if train_missing or test_missing:
        log_lines.append(
            f"Note: {train_missing} train / {test_missing} test images referenced in the "
            "JSON were not found among the uploaded files and were skipped."
        )

    label_map = fu.build_label_map(train_cats, test_cats)
    label_names = sorted(label_map, key=label_map.get)

    if not label_map:
        yield "No usable categories found in the COCO annotations. Check your export.", None, None
        return

    log_lines.append(f"Found {len(train_samples)} train / {len(test_samples)} test images.")
    log_lines.append(f"Classes: {label_names}")
    if HAS_ZEROGPU:
        log_lines.append("ZeroGPU detected — this run will use a real GPU.")
    yield "\n".join(log_lines), None, None

    input_size = int(input_size)
    batch_size = int(batch_size)
    grid_size = input_size // 8

    train_ds = fu.make_dataset(
        train_samples, label_map, input_size, grid_size, batch_size,
        augment=True, cat_id_to_name=train_cats,
    )
    val_ds = fu.make_dataset(
        test_samples, label_map, input_size, grid_size, batch_size,
        augment=False, cat_id_to_name=test_cats,
    )
    steps_per_epoch = max(1, math.ceil(len(train_samples) / batch_size))
    validation_steps = max(1, math.ceil(len(test_samples) / batch_size))

    log_lines.append("Building model (MobileNetV2 backbone, alpha=0.35)...")
    yield "\n".join(log_lines), None, None
    model = fu.build_fomo_model(input_size, num_classes=len(label_map))

    log_lines.append(f"Training with {optimizer_name.upper()}...")
    yield "\n".join(log_lines), None, None

    epoch_logs = []
    history = fu.train_model(
        model, train_ds, val_ds, int(epochs), float(learning_rate),
        float(background_weight), epoch_logs, optimizer_name=optimizer_name,
        steps_per_epoch=steps_per_epoch, validation_steps=validation_steps,
    )
    log_lines.extend(epoch_logs)
    yield "\n".join(log_lines), None, None

    fig, ax = plt.subplots()
    ax.plot(history.history["loss"], label="train loss")
    ax.plot(history.history["val_loss"], label="val loss")
    ax.set_xlabel("epoch")
    ax.set_ylabel("loss")
    ax.legend()
    fig.savefig(PLOT_PATH)
    plt.close(fig)

    log_lines.append("Quantizing to int8 TFLite...")
    yield "\n".join(log_lines), PLOT_PATH, None

    tflite_bytes = fu.quantize_model(model, train_samples, input_size)
    with open(TFLITE_PATH, "wb") as f:
        f.write(tflite_bytes)

    h_text = fu.tflite_to_c_header(tflite_bytes)
    with open(H_PATH, "w") as f:
        f.write(h_text)

    with open(LABELS_PATH, "w") as f:
        json.dump(label_names, f)

    size_kb = len(tflite_bytes) / 1024
    log_lines.append(f"Done. Quantized model size: {size_kb:.1f} KB")
    log_lines.append("Go to the 'Test Model' tab to try it, or 'Download' to grab the files.")
    yield "\n".join(log_lines), PLOT_PATH, TFLITE_PATH


def run_test(image, threshold):
    if not os.path.exists(TFLITE_PATH) or not os.path.exists(LABELS_PATH):
        return None, "Train a model first (see the Train tab)."
    if image is None:
        return None, "Upload an image to test."

    with open(LABELS_PATH, "r") as f:
        label_names = json.load(f)

    draw_img, detections = fu.run_tflite_inference(
        TFLITE_PATH, image, label_names, threshold=threshold,
    )
    if detections:
        summary = f"{len(detections)} detection(s):\n" + "\n".join(
            f"- {d['class']} ({d['confidence']:.2f}) at ({d['x']:.0f}, {d['y']:.0f})"
            for d in detections
        )
    else:
        summary = "No detections above threshold."
    return draw_img, summary


def get_download_files():
    if not os.path.exists(TFLITE_PATH):
        return None, None, None
    return TFLITE_PATH, H_PATH, LABELS_PATH


with gr.Blocks(title="Crow FOMO Trainer") as demo:
    gr.Markdown(
        "# Crow Detector — FOMO-style TFLite Trainer\n"
        "Train a lightweight object detector, quantize it to int8, "
        "and export it for the Seeed XIAO ESP32S3 Sense.\n\n"
        "**Before you start:** export your Roboflow dataset in **COCO JSON** format. "
        "You'll get a `train` folder and a `test` (or `valid`) folder, each containing "
        "images plus `_annotations.coco.json`. Click each box below and select the "
        "whole folder (not individual files)."
    )

    with gr.Tab("1. Train"):
        with gr.Row():
            train_files = gr.File(
                label="Train folder (images + _annotations.coco.json)",
                file_count="directory",
            )
            test_files = gr.File(
                label="Test folder (images + _annotations.coco.json)",
                file_count="directory",
            )
        with gr.Row():
            input_size = gr.Dropdown(["96", "112", "160"], value="96", label="Input size (px)")
            epochs = gr.Slider(5, 100, value=30, step=1, label="Epochs")
            batch_size = gr.Slider(4, 64, value=16, step=4, label="Batch size")
        with gr.Row():
            learning_rate = gr.Number(value=0.001, label="Learning rate")
            background_weight = gr.Slider(
                0.05, 1.0, value=0.3, step=0.05,
                label="Background loss weight (lower = focus more on crows, helps with class imbalance)",
            )
            optimizer_name = gr.Dropdown(
                ["adamw", "adam", "sgd"], value="adamw", label="Optimizer",
            )
        train_btn = gr.Button("Start Training", variant="primary")
        train_log = gr.Textbox(label="Training log", lines=15)
        loss_plot = gr.Image(label="Loss curve")
        tflite_out = gr.File(label="Quantized model (also available in Download tab)")

        train_btn.click(
            run_training,
            inputs=[train_files, test_files, input_size, epochs, batch_size, learning_rate, background_weight, optimizer_name],
            outputs=[train_log, loss_plot, tflite_out],
        )

    with gr.Tab("2. Test Model"):
        with gr.Row():
            test_image = gr.Image(label="Upload a test image", type="pil")
            result_image = gr.Image(label="Detections")
        threshold = gr.Slider(0.1, 0.95, value=0.5, step=0.05, label="Confidence threshold")
        test_btn = gr.Button("Run Detection")
        result_text = gr.Textbox(label="Results")
        test_btn.click(run_test, inputs=[test_image, threshold], outputs=[result_image, result_text])

    with gr.Tab("3. Download"):
        gr.Markdown(
            "Download the quantized model for deployment.\n\n"
            "- **.tflite** — for use with `esp-tflite-micro` if you load the model from flash/LittleFS\n"
            "- **.h** — the same model as a C byte array, ready to `#include` directly in an Arduino sketch\n"
            "- **labels.json** — class index to name mapping"
        )
        refresh_btn = gr.Button("Refresh available files")
        tflite_file = gr.File(label="crow_fomo_int8.tflite")
        h_file = gr.File(label="crow_fomo_model.h")
        labels_file = gr.File(label="labels.json")
        refresh_btn.click(get_download_files, outputs=[tflite_file, h_file, labels_file])

if __name__ == "__main__":
    demo.launch()