Create app.py

app.py

@@ -0,0 +1,204 @@
+# app.py
+import io
+import os
+from typing import List, Tuple
+import requests
+from PIL import Image, UnidentifiedImageError
+import pandas as pd
+
+import gradio as gr
+from transformers import pipeline
+
+# Model choices
+MODEL_CHOICES = {
+    "shadowlilac/aesthetic-shadow (v1, fp32)": {
+        "repo": "shadowlilac/aesthetic-shadow",
+        "precision": "fp32",
+    },
+    "NeoChen1024/aesthetic-shadow-v2-backup (fp32)": {
+        "repo": "NeoChen1024/aesthetic-shadow-v2-backup",
+        "precision": "fp32",
+    },
+    "Disty0/aesthetic-shadow-v2 (fp16)": {
+        "repo": "Disty0/aesthetic-shadow-v2",
+        "precision": "fp16",
+    },
+    "default": "Disty0/aesthetic-shadow-v2 (fp16)"
+}
+
+# Keep a global reference to the current pipeline
+pipe = None
+current_model_repo = None
+
+
+def load_model(model_key: str):
+    """Load a model by dropdown key if not already loaded."""
+    global pipe, current_model_repo
+    info = MODEL_CHOICES[model_key]
+    repo = info["repo"]
+    if repo == current_model_repo and pipe is not None:
+        return pipe
+    # Load new pipeline on CPU
+    pipe = pipeline(
+        "image-classification",
+        model=repo,
+        device=-1,
+    )
+    current_model_repo = repo
+    return pipe
+
+
+def pil_from_uploaded(uploaded) -> Image.Image:
+    if uploaded is None:
+        return None
+    if hasattr(uploaded, "name"):
+        try:
+            return Image.open(uploaded).convert("RGB")
+        except UnidentifiedImageError:
+            return None
+    if isinstance(uploaded, Image.Image):
+        return uploaded.convert("RGB")
+    return None
+
+
+def pil_from_url(url: str) -> Image.Image:
+    if not url:
+        return None
+    try:
+        r = requests.get(url, timeout=10)
+        r.raise_for_status()
+        return Image.open(io.BytesIO(r.content)).convert("RGB")
+    except Exception:
+        return None
+
+
+def extract_hq_score(preds) -> float:
+    for p in preds:
+        if str(p.get("label")).lower() == "hq":
+            return float(p.get("score", 0.0))
+    if len(preds):
+        return float(preds[0].get("score", 0.0))
+    return 0.0
+
+
+def make_progress_html(score: float) -> str:
+    pct = int(round(score * 100))
+    return f"""
+    <div style="width:100%; border:1px solid #ddd; border-radius:6px; padding:6px;">
+      <div style="font-weight:600; margin-bottom:4px;">High-quality score: {score:.3f} ({pct}%)</div>
+      <div style="background:#eee; border-radius:4px; overflow:hidden;">
+        <div style="width:{pct}%; padding:6px 0; text-align:center; font-weight:600;">
+          {pct}%
+        </div>
+      </div>
+    </div>
+    """
+
+
+def classify_images(images: List[Image.Image], pipe) -> List[float]:
+    if not images:
+        return []
+    results = pipe(images=images)
+    scores = [extract_hq_score(r) for r in results]
+    return scores
+
+
+def run_classify(
+    uploaded_image,
+    url_input,
+    batch_files,
+    batch_urls_text,
+    model_key,
+) -> Tuple[str, List[List], dict]:
+    images_for_batch = []
+    names_for_batch = []
+
+    if batch_files:
+        for f in batch_files:
+            img = pil_from_uploaded(f)
+            if img is not None:
+                images_for_batch.append(img)
+                name = getattr(f, "name", "uploaded_file")
+                names_for_batch.append(os.path.basename(name))
+    if batch_urls_text:
+        for line in batch_urls_text.splitlines():
+            line = line.strip()
+            if not line:
+                continue
+            img = pil_from_url(line)
+            if img is not None:
+                images_for_batch.append(img)
+                names_for_batch.append(line)
+
+    pipe = load_model(model_key)
+
+    if images_for_batch:
+        scores = classify_images(images_for_batch, pipe)
+        rows = [[names_for_batch[i] if i < len(names_for_batch) else f"img_{i}", float(scores[i])] for i in range(len(scores))]
+        avg = sum(scores) / len(scores)
+        html = make_progress_html(avg)
+        return html, rows, {"mode": "batch"}
+
+    img = None
+    img_name = "input"
+    if url_input:
+        img = pil_from_url(url_input.strip())
+        img_name = url_input.strip()
+    if img is None and uploaded_image:
+        img = pil_from_uploaded(uploaded_image)
+        img_name = getattr(uploaded_image, "name", "uploaded_image")
+
+    if img is None:
+        return "<div style='color:#a00;'>No valid image(s) provided. Please upload or supply a URL.</div>", [], {"mode": "none"}
+
+    scores = classify_images([img], pipe)
+    score = float(scores[0]) if scores else 0.0
+    html = make_progress_html(score)
+    rows = [[img_name, score]]
+    return html, rows, {"mode": "single", "image": img}
+
+
+# Build the Gradio UI
+with gr.Blocks(title="Aesthetic Shadow - Anime Image Quality Classifier (CPU)") as demo:
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            with gr.Tabs():
+                with gr.TabItem("Image Upload"):
+                    uploaded_image = gr.File(label="Upload single image", file_count="single", type="file")
+                with gr.TabItem("URL"):
+                    url_input = gr.Textbox(label="Image URL", placeholder="https://...", lines=1)
+                with gr.TabItem("Batch"):
+                    batch_files = gr.File(label="Upload multiple images (batch)", file_count="multiple", type="file")
+                    batch_urls_text = gr.Textbox(label="Batch URLs (one per line)", placeholder="https://...", lines=4)
+            gr.Markdown("- If batch inputs are provided they will be used as batch mode. Otherwise single image/url will be used.")
+
+        with gr.Column(scale=1):
+            model_dropdown = gr.Dropdown(
+                choices=[k for k in MODEL_CHOICES.keys() if k != "default"],
+                value=MODEL_CHOICES["default"],
+                label="Model Selection",
+            )
+            run_button = gr.Button("Run", variant="primary")
+
+            result_html = gr.HTML("<div>Result will appear here after running.</div>")
+            result_image = gr.Image(label="Input image (single mode)", interactive=False)
+            result_table = gr.Dataframe(headers=["source", "hq_score"], interactive=False)
+
+            def run_handler(up_img, url_txt, b_files, b_urls, model_key):
+                html, rows, meta = run_classify(up_img, url_txt, b_files, b_urls, model_key)
+                if meta.get("mode") == "single" and meta.get("image") is not None:
+                    return html, meta.get("image"), pd.DataFrame(rows, columns=["source", "hq_score"])
+                else:
+                    return html, None, pd.DataFrame(rows, columns=["source", "hq_score"])
+
+            run_button.click(
+                fn=run_handler,
+                inputs=[uploaded_image, url_input, batch_files, batch_urls_text, model_dropdown],
+                outputs=[result_html, result_image, result_table],
+            )
+
+    gr.Markdown("All Aesthetic Shadow models are by shadowlilac. V2 is using reuploads by other people.")
+
+if __name__ == "__main__":
+    demo.launch()