docker/app.py

"""Docker Gradio app for the ai-image-detector model.

Loads the merged CLIP ViT-B/16 LoRA weights from ``model.safetensors`` (next to
this file, or under the directory pointed to by ``MODEL_DIR``) and exposes a
Gradio interface for uploading an image and getting a real/fake/uncertain
prediction. Mirrors the Hugging Face Space ``app.py`` but is parametrised so it
can run from a container that downloads the model at runtime.
"""
from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

import gradio as gr
import torch
from PIL import Image

APP_CSS = """
.gradio-container { max-width: 1040px !important; }
.detector-header { align-items: center; display: flex; justify-content: space-between; margin-bottom: 16px; }
.detector-title { font-size: 26px; font-weight: 700; line-height: 1.15; }
.detector-meta { color: var(--body-text-color-subdued); font-size: 13px; text-align: right; }
.result-card { background: var(--background-fill-secondary); border-radius: 8px; padding: 14px 16px; }
.result-heading { display: flex; gap: 12px; justify-content: space-between; margin-bottom: 10px; }
.result-label { font-size: 24px; font-weight: 800; line-height: 1.1; }
.result-pill { border-radius: 999px; color: white; font-size: 12px; font-weight: 800; height: fit-content; padding: 5px 10px; }
.result-real { background: #15803d; }
.result-fake { background: #b91c1c; }
.result-uncertain { background: #b45309; }
.metric-grid { display: grid; gap: 12px; grid-template-columns: repeat(auto-fit, minmax(120px, 1fr)); margin-top: 8px; }
.metric { display: flex; flex-direction: column; gap: 2px; }
.metric-name { color: var(--body-text-color-subdued); font-size: 12px; }
.metric-value { font-size: 18px; font-weight: 700; }
"""


def _resolve_model_dir() -> Path:
    """Resolve where model.safetensors + config.json live.

    Priority: MODEL_DIR env var > directory of this file.
    """
    env_dir = os.environ.get("MODEL_DIR")
    if env_dir:
        return Path(env_dir).expanduser().resolve()
    return Path(__file__).resolve().parent


def _load_model(model_dir: Path):
    """Load merged weights + config from ``model_dir``."""
    import timm
    from safetensors.torch import load_file

    weights_path = model_dir / "model.safetensors"
    config_path = model_dir / "config.json"
    if not weights_path.exists():
        raise FileNotFoundError(
            f"Model weights not found at {weights_path}. "
            "Set MODEL_DIR or run the download_model.py entrypoint first."
        )
    if not config_path.exists():
        raise FileNotFoundError(f"config.json not found at {config_path}.")
    cfg = json.loads(config_path.read_text())
    model = timm.create_model(
        cfg["backbone"], pretrained=False, num_classes=1, img_size=cfg["image_size"]
    )
    state = load_file(str(weights_path))
    missing, unexpected = model.load_state_dict(state, strict=False)
    if unexpected:
        print(f"[app] Unexpected keys in safetensors (ignored): {len(unexpected)}")
    if missing:
        print(f"[app] Missing keys when loading weights: {len(missing)}")
    model.eval()
    return model, cfg


class DockerPredictor:
    """Lightweight predictor that loads the merged model from safetensors."""

    def __init__(self, model_dir: Path | None = None):
        model_dir = model_dir or _resolve_model_dir()
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"[app] Loading model from {model_dir} on {self.device}")
        self.model, self.cfg = _load_model(model_dir)
        self.model = self.model.to(self.device)
        self.real_threshold = float(self.cfg.get("real_threshold", 0.93))
        self.fake_threshold = float(self.cfg.get("fake_threshold", 0.91))
        self.temperature = float(self.cfg.get("temperature", 1.0))
        self.img_size = int(self.cfg.get("image_size", 256))
        mean = self.cfg.get("normalization_mean", [0.481, 0.458, 0.408])
        std = self.cfg.get("normalization_std", [0.269, 0.261, 0.276])
        from torchvision import transforms

        self.transform = transforms.Compose(
            [
                transforms.Resize((self.img_size, self.img_size)),
                transforms.ToTensor(),
                transforms.Normalize(mean=mean, std=std),
            ]
        )

    @torch.inference_mode()
    def predict(self, image: Image.Image) -> dict[str, Any]:
        if image is None:
            return {}
        image = image.convert("RGB")
        x = self.transform(image).unsqueeze(0).to(self.device)
        logit = self.model(x).reshape(-1)
        if self.temperature and self.temperature > 0:
            logit = logit / self.temperature
        p_real = float(torch.sigmoid(logit).item())
        if p_real < self.fake_threshold:
            prediction, confidence = "fake", 1.0 - p_real
        elif p_real >= self.real_threshold:
            prediction, confidence = "real", p_real
        else:
            prediction, confidence = "uncertain", max(
                p_real - self.fake_threshold, self.real_threshold - p_real
            )
        return {
            "prediction": prediction,
            "confidence": confidence,
            "real_probability": p_real,
            "fake_probability": 1.0 - p_real,
            "fake_threshold": self.fake_threshold,
            "real_threshold": self.real_threshold,
            "img_size": self.img_size,
            "temperature": self.temperature,
        }


def _result_card(m: dict[str, Any]) -> str:
    prediction = m.get("prediction", "unknown")
    pill = {"real": "result-real", "fake": "result-fake"}.get(prediction, "result-uncertain")
    return f"""
<div class="result-card">
  <div class="result-heading">
    <div>
      <div class="metric-name">Prediction</div>
      <div class="result-label">{prediction.upper()}</div>
    </div>
    <div class="result-pill {pill}">{prediction.upper()}</div>
  </div>
  <div class="metric-grid">
    <div class="metric"><div class="metric-name">Confidence</div><div class="metric-value">{m.get('confidence', 0):.4f}</div></div>
    <div class="metric"><div class="metric-name">p(real)</div><div class="metric-value">{m.get('real_probability', 0):.4f}</div></div>
    <div class="metric"><div class="metric-name">p(fake)</div><div class="metric-value">{m.get('fake_probability', 0):.4f}</div></div>
    <div class="metric"><div class="metric-name">Fake threshold</div><div class="metric-value">{m.get('fake_threshold', 0)}</div></div>
    <div class="metric"><div class="metric-name">Real threshold</div><div class="metric-value">{m.get('real_threshold', 0)}</div></div>
  </div>
</div>
"""


def _make_predict_fn(predictor: DockerPredictor):
    """Bind the predictor into a single-arg Gradio handler."""

    def _predict(image: Image.Image | None):
        if image is None:
            return '<div class="result-card">Lütfen bir görsel yükleyin.</div>', {}
        m = predictor.predict(image)
        return _result_card(m), m

    return _predict


def build_demo(predictor: DockerPredictor, examples_dir: Path | None = None) -> gr.Blocks:
    predict_fn = _make_predict_fn(predictor)
    with gr.Blocks(title="AI Image Detector", css=APP_CSS) as demo:
        gr.HTML(
            f"""
<div class="detector-header">
  <div class="detector-title">AI Image Detector</div>
  <div class="detector-meta">threshold {predictor.real_threshold:.2f}<br>CLIP ViT-B/16 + LoRA</div>
</div>
"""
        )
        with gr.Tab("Tek görsel"):
            with gr.Row():
                with gr.Column():
                    img_in = gr.Image(type="pil", label="Görsel")
                    btn = gr.Button("Tahmin et")
                with gr.Column():
                    out_card = gr.HTML(label="Sonuç")
                    raw = gr.Json(label="Ham skorlar")
            btn.click(predict_fn, inputs=[img_in], outputs=[out_card, raw])
            if examples_dir and examples_dir.exists():
                examples = sorted(examples_dir.glob("*.jpg"))
                if examples:
                    gr.Examples(
                        examples=[[str(p)] for p in examples],
                        inputs=[img_in],
                        outputs=[out_card, raw],
                        fn=predict_fn,
                        cache_examples=False,
                    )
    return demo


def main() -> int:
    model_dir = _resolve_model_dir()
    examples_dir = model_dir / "examples"
    predictor = DockerPredictor(model_dir=model_dir)
    server_name = os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0")
    server_port = int(os.environ.get("GRADIO_SERVER_PORT", "7860"))
    demo = build_demo(predictor, examples_dir=examples_dir)
    # share=False: we never want a public tunnel from inside a container.
    demo.launch(
        server_name=server_name,
        server_port=server_port,
        share=False,
        show_error=True,
        inbrowser=False,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())