| """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) |
| |
| 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()) |
|
|
