Update app.py

app.py

@@ -1,19 +1,17 @@
 import os
 import gradio as gr
-import json
-from typing import Any, Dict, Tuple
 from dotenv import load_dotenv
-from urllib.request import urlopen, Request, urlretrieve
-from io import BytesIO
-from PIL import Image
-from functools import lru_cache
+from urllib.request import urlretrieve
+from utils import (
+    EXAMPLE_ITEMS,
+    analyze_image_with_status,
+    run_example_by_index,
+)
 
 load_dotenv()
 
 VIDDEXA_TOKEN = os.getenv("HF_TOKEN")
 
-_MODEL_CACHE: Dict[str, Any] = {}
-
 CUSTOM_CSS = """
 .verdict-safe { background-color: #D5F5E3; border: 2px solid #2ECC71; color: #1D8348; }
 .verdict-sensitive { background-color: #FCF3CF; border: 2px solid #F1C40F; color: #B7950B; }
@@ -23,162 +21,6 @@ CUSTOM_CSS = """
 footer {visibility: hidden}
 """
 
-
-@lru_cache(maxsize=32)
-def _download_image_bytes(image_url: str) -> bytes:
-    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) -> Any:
-    """Loads a model and caches 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=VIDDEXA_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:
-    """Gets 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]:
-    """Formats 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,
-                  progress=gr.Progress(track_tqdm=True)):
-    """The 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)
-    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,
-                              progress=gr.Progress(track_tqdm=True)):
-    """Run analysis and also return a user-friendly status string."""
-    verdict_html, label_scores, md_scores, json_obj = analyze_image(image_path, image_url, model_choice, 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
-
-
-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)",
-    ),
-]
-
-
-def run_example_by_index(evt: gr.SelectData):
-    """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)
-    status = f"Last analysed example: {caption}"
-    return (
-        verdict_html,
-        label_scores,
-        md_scores,
-        json_obj,
-        gr.update(value=model),
-        gr.update(value=url),
-        status,
-    )
-
 with gr.Blocks(
     theme=gr.themes.Default(primary_hue=gr.themes.colors.teal),
     title="Visual Content Moderation",
@@ -261,7 +103,7 @@ with gr.Blocks(
     status_md = gr.Markdown("Last analysed: —", elem_id="last-example-status")
 
     examples_gallery.select(
-        fn=run_example_by_index,
+        fn=lambda evt: run_example_by_index(evt, VIDDEXA_TOKEN),
         outputs=[
             verdict_output,
             label_output,
@@ -274,7 +116,7 @@ with gr.Blocks(
     )
 
     run_btn.click(
-        fn=analyze_image_with_status,
+        fn=lambda img, url, model, progress: analyze_image_with_status(img, url, model, VIDDEXA_TOKEN, progress),
         inputs=[image_input, image_url_input, model_choice],
         outputs=[verdict_output, label_output, markdown_output, json_output, status_md],
     )