"""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"""
Prediction
{prediction.upper()}
{prediction.upper()}
Confidence
{m.get('confidence', 0):.4f}
p(real)
{m.get('real_probability', 0):.4f}
p(fake)
{m.get('fake_probability', 0):.4f}
Fake threshold
{m.get('fake_threshold', 0)}
Real threshold
{m.get('real_threshold', 0)}
"""
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 '
Lütfen bir görsel yükleyin.
', {}
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"""