Create utils.py

utils.py

@@ -0,0 +1,166 @@
+import os
+import json
+import gradio as gr
+from typing import Any, Dict, Tuple
+from urllib.request import urlopen, Request
+from io import BytesIO
+from PIL import Image
+from functools import lru_cache
+
+_MODEL_CACHE: Dict[str, Any] = {}
+
+EXAMPLE_ITEMS = [
+    (
+        "https://assets.clevelandclinic.org/transform/LargeFeatureImage/cd71f4bd-81d4-45d8-a450-74df78e4477a/Apples-184940975-770x533-1_jpg",
+        "viddexa/nsfw-mini",
+        "Apples (mini)",
+    ),
+    (
+        "https://img.freepik.com/free-photo/breast-screening-is-very-important-every-woman_329181-14953.jpg",
+        "viddexa/nsfw-nano",
+        "Breast screening (nano)",
+    ),
+    (
+        "https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcSbRwt56NYsiHwrT8oS-igzgeEzp7p3Jbe2dw&s",
+        "viddexa/nsfw-mini",
+        "Thumbnail (mini)",
+    ),
+    (
+        "https://img.freepik.com/premium-photo/portrait-beautiful-young-woman_1048944-5548042.jpg",
+        "viddexa/nsfw-nano",
+        "Portrait (nano)",
+    ),
+]
+
+
+@lru_cache(maxsize=32)
+def _download_image_bytes(image_url: str) -> bytes:
+    """Download image bytes from URL with caching."""
+    req = Request(image_url, headers={"User-Agent": "viddexa-gradio-demo/1.0"})
+    with urlopen(req, timeout=20) as resp:
+        return resp.read()
+
+
+def _load_model(model_id: str, token: str | None = None) -> Any:
+    """Load a model and cache it."""
+    if model_id in _MODEL_CACHE:
+        return _MODEL_CACHE[model_id]
+    try:
+        from moderators.auto_model import AutoModerator
+        model = AutoModerator.from_pretrained(model_id, token=token, use_fast=True)
+        _MODEL_CACHE[model_id] = model
+        return model
+    except Exception as e:
+        error_msg = f"Failed to load model: {model_id}. Error: {e}"
+        if "401" in str(e):
+            error_msg += "\n\nThis model may be private. Please ensure you have provided a valid Hugging Face token if required."
+        raise gr.Error(error_msg)
+
+
+def _get_image_input(image_path: str | None, image_url: str | None) -> Image.Image:
+    """Get image data from either an uploaded file path or a URL."""
+    if image_url:
+        try:
+            data = _download_image_bytes(image_url)
+            img = Image.open(BytesIO(data))
+            return img.convert("RGB")
+        except Exception as fetch_err:
+            raise gr.Error(f"Could not download or open the image from the URL: {fetch_err}")
+    elif image_path:
+        img = Image.open(image_path)
+        return img.convert("RGB")
+    else:
+        raise gr.Error("Please upload an image or provide an image URL.")
+
+
+def _format_results(results: list) -> Tuple[str, Dict[str, float], str, Dict]:
+    """Format the model output for the Gradio interface."""
+    if not results or "classifications" not in results[0]:
+        return "<div class='verdict-card'>No classifications found.</div>", {}, "No classifications found.", {}
+
+    classifications = results[0]["classifications"]
+
+    label_output: Dict[str, float]
+    if isinstance(classifications, dict):
+        label_output = {str(k): float(v) for k, v in classifications.items()}
+    else:
+        try:
+            label_output = {str(item['label']): float(item['score']) for item in classifications}
+        except Exception:
+            label_output = {}
+
+    scores = {label.lower(): score for label, score in label_output.items()}
+    nsfw_score = scores.get("nsfw", 0.0)
+
+    if nsfw_score > 0.7:
+        verdict_text = "HIGH RISK: NSFW"
+        verdict_class = "verdict-nsfw"
+    elif nsfw_score > 0.2:
+        verdict_text = "MEDIUM RISK: SENSITIVE"
+        verdict_class = "verdict-sensitive"
+    else:
+        verdict_text = "LOW RISK: SAFE"
+        verdict_class = "verdict-safe"
+
+    verdict_html = f"<div class='verdict-card {verdict_class}'>{verdict_text}</div>"
+
+    markdown_output = "### All Scores\n---\n"
+    for label, score in sorted(label_output.items(), key=lambda kv: kv[1], reverse=True):
+        markdown_output += f"- **{label.capitalize()}**: {score:.4f}\n"
+
+    return verdict_html, label_output, markdown_output, results[0]
+
+
+def analyze_image(image_path: str | None, image_url: str | None, model_choice: str,
+                  token: str | None = None, progress=gr.Progress(track_tqdm=True)):
+    """Main inference function for the Gradio interface."""
+    progress(0, desc="Initializing Analysis...")
+    progress(0.2, desc="Processing Image...")
+    input_image = _get_image_input(image_path, image_url)
+    progress(0.5, desc=f"Loading Model: {os.path.basename(model_choice)}...")
+    model = _load_model(model_choice, token)
+    progress(0.8, desc="Running Inference...")
+    results = model(input_image)
+
+    json_results = [
+        {"classifications": getattr(r, "classifications", r)}
+        for r in results
+    ]
+    json_results = json.loads(json.dumps(json_results, ensure_ascii=False))
+
+    progress(1, desc="Complete!")
+    return _format_results(json_results)
+
+
+def analyze_image_with_status(image_path: str | None, image_url: str | None, model_choice: str,
+                              token: str | None = None, progress=gr.Progress(track_tqdm=True)):
+    """Run analysis and return results with user-friendly status string."""
+    verdict_html, label_scores, md_scores, json_obj = analyze_image(image_path, image_url, model_choice, token, progress)
+    if image_url:
+        status = f"Last analysed URL: {image_url}"
+    elif image_path:
+        status = "Last analysed uploaded image."
+    else:
+        status = "Last analysed: —"
+    return verdict_html, label_scores, md_scores, json_obj, status
+
+
+def run_example_by_index(evt: gr.SelectData, token: str | None = None):
+    """Handle gallery selection: run analysis for the selected example and update inputs."""
+    try:
+        idx = int(getattr(evt, "index", 0))
+    except Exception:
+        idx = 0
+    idx = max(0, min(idx, len(EXAMPLE_ITEMS) - 1))
+    url, model, caption = EXAMPLE_ITEMS[idx]
+    verdict_html, label_scores, md_scores, json_obj = analyze_image(None, url, model, token)
+    status = f"Last analysed example: {caption}"
+    return (
+        verdict_html,
+        label_scores,
+        md_scores,
+        json_obj,
+        gr.update(value=model),
+        gr.update(value=url),
+        status,
+    )