refactor: simplify utils.py

utils.py

@@ -1,7 +1,5 @@
-import os
-import json
 import gradio as gr
-from typing import Any, Dict, Tuple
+from typing import Any, Dict
 from urllib.request import urlopen, Request
 from io import BytesIO
 from PIL import Image
@@ -10,157 +8,44 @@ 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)",
-    ),
+    ("https://assets.clevelandclinic.org/transform/LargeFeatureImage/cd71f4bd-81d4-45d8-a450-74df78e4477a/Apples-184940975-770x533-1_jpg", "viddexa/nsfw-detection-mini"),
+    ("https://img.freepik.com/free-photo/breast-screening-is-very-important-every-woman_329181-14953.jpg", "viddexa/nsfw-detection-nano"),
+    ("https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcSbRwt56NYsiHwrT8oS-igzgeEzp7p3Jbe2dw&s", "viddexa/nsfw-detection-mini"),
+    ("https://img.freepik.com/premium-photo/portrait-beautiful-young-woman_1048944-5548042.jpg", "viddexa/nsfw-detection-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"})
+def download_image(url: str) -> Image.Image:
+    """Download and return PIL Image from URL."""
+    req = Request(url, headers={"User-Agent": "viddexa-gradio-demo/1.0"})
     with urlopen(req, timeout=20) as resp:
-        return resp.read()
+        return Image.open(BytesIO(resp.read())).convert("RGB")
 
 
-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:
+def load_model(model_id: str, token: str | None = None) -> Any:
+    """Load model with caching."""
+    if model_id not in _MODEL_CACHE:
         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,
-    )
+        _MODEL_CACHE[model_id] = AutoModerator.from_pretrained(model_id, token=token, use_fast=True)
+    return _MODEL_CACHE[model_id]
+
+
+def analyze(image_path: str | None, image_url: str | None, model_id: str, token: str | None = None):
+    """Run inference and return classification scores."""
+    if not image_url and not image_path:
+        raise gr.Error("Provide an image or URL")
+    
+    img = download_image(image_url) if image_url else Image.open(image_path).convert("RGB")
+    model = load_model(model_id, token)
+    results = model(img)
+    
+    classifications = results[0].classifications if hasattr(results[0], "classifications") else results[0]["classifications"]
+    return {str(k): float(v) for k, v in (classifications.items() if isinstance(classifications, dict) else [(c["label"], c["score"]) for c in classifications])}
+
+
+def run_example(evt: gr.SelectData, token: str | None = None):
+    """Handle example selection."""
+    idx = max(0, min(int(getattr(evt, "index", 0)), len(EXAMPLE_ITEMS) - 1))
+    url, model = EXAMPLE_ITEMS[idx]
+    return analyze(None, url, model, token), gr.update(value=model), gr.update(value=url)