Update app.py

app.py

@@ -1,38 +1,20 @@
-import os, shutil, zipfile, pathlib, traceback, math
+import os, zipfile, pathlib, shutil, traceback
 import numpy as np
 import pandas as pd
-
-import gradio as gr
 from PIL import Image, ImageOps
-
-from datasets import load_dataset
+import gradio as gr
 from huggingface_hub import hf_hub_download
-
-# AutoGluon (multimodal)
 from autogluon.multimodal import MultiModalPredictor
+from datasets import load_dataset
 
-# ---------------- Settings ----------------
-TITLE = "💄 Lipstick Detection (EfficientNet-B0 via AutoGluon)"
-DESC  = (
-    "- Upload a face image; the model predicts **Lipstick** vs **No Lipstick**.\n"
-    "- Left: original; Right: the **preprocessed** 224×224 view used by the model.\n"
-    "- Teaching demo only."
-)
-
-# =========================
-# CONFIG
-# =========================
 MODEL_REPO_ID = "Iris314/nn_automl_model"
 CANDIDATE_FILES = ["best_model.zip", "best_model.pkl"]
-DATASET_REPO   = "keerthikoganti/lipstick-image-dataset"
+DATASET_REPO = "keerthikoganti/lipstick-image-dataset"
 
-CACHE_DIR  = pathlib.Path("hf_cache")
-MODEL_DIR  = CACHE_DIR / "model"
+CACHE_DIR = pathlib.Path("hf_cache")
+MODEL_DIR = CACHE_DIR / "model"
 CACHE_DIR.mkdir(exist_ok=True, parents=True)
 
-# =========================
-# Load trained predictor
-# =========================
 def load_mm_predictor():
     downloaded = None
     for fname in CANDIDATE_FILES:
@@ -49,9 +31,8 @@ def load_mm_predictor():
         except Exception:
             pass
     if downloaded is None:
-        raise FileNotFoundError(f"Could not find any of {CANDIDATE_FILES} in {MODEL_REPO_ID} ")
+        raise FileNotFoundError(f"Could not find any of {CANDIDATE_FILES} in {MODEL_REPO_ID}")
 
-    # Prepare folder for MultiModalPredictor.load
     if MODEL_DIR.exists():
         shutil.rmtree(MODEL_DIR)
     MODEL_DIR.mkdir(parents=True, exist_ok=True)
@@ -60,98 +41,51 @@ def load_mm_predictor():
         with zipfile.ZipFile(downloaded, "r") as zf:
             zf.extractall(MODEL_DIR)
         kids = [p for p in MODEL_DIR.iterdir()]
-        load_path = MODEL_DIR
-        if len(kids) == 1 and kids[0].is_dir():
-            load_path = kids[0]
+        load_path = kids[0] if len(kids) == 1 and kids[0].is_dir() else MODEL_DIR
     else:
         load_path = downloaded
 
     predictor = MultiModalPredictor.load(str(load_path))
     return predictor
 
-try:
-    PREDICTOR = load_mm_predictor()
-except Exception as e:
-    PREDICTOR = None
-    print("Failed to load predictor:", e)
+PREDICTOR = load_mm_predictor()
 
-# =========================
-# Preprocess (for visualization)
-# EfficientNet-B0 typical viz: Resize(shorter side=256) → CenterCrop(224)
-# =========================
 TARGET_SIZE = 224
 RESIZE_SHORT = 256
 
 def preprocess_for_viz(pil_img: Image.Image) -> Image.Image:
     img = pil_img.convert("RGB")
-    img = ImageOps.exif_transpose(img)  # respect orientation EXIF
+    img = ImageOps.exif_transpose(img)
     w, h = img.size
-    if w <= 0 or h <= 0:
-        raise ValueError("Invalid image dimensions.")
     scale = RESIZE_SHORT / min(w, h)
     new_w, new_h = int(round(w * scale)), int(round(h * scale))
     img = img.resize((new_w, new_h), Image.BICUBIC)
     left = (new_w - TARGET_SIZE) // 2
-    top  = (new_h - TARGET_SIZE) // 2
+    top = (new_h - TARGET_SIZE) // 2
     img = img.crop((left, top, left + TARGET_SIZE, top + TARGET_SIZE))
     return img
 
-# =========================
-# File validation
-# =========================
-ALLOWED_EXT = {".jpg", ".jpeg", ".png"}
-MAX_BYTES   = 8 * 1024 * 1024  # 8 MB
-MAX_SIDE    = 4096
-
-def validate_image(fileobj) -> str:
-    if fileobj is None:
-        raise ValueError("Please upload an image.")
-    # gr.Image with type="filepath" returns a str path in Spaces
-    path = getattr(fileobj, "name", fileobj)
-    if not isinstance(path, str):
-        path = str(path)
-
-    ext = pathlib.Path(path).suffix.lower()
-    if ext not in ALLOWED_EXT:
-        raise ValueError("Unsupported file type. Please upload a PNG or JPG/JPEG.")
-    if os.path.getsize(path) > MAX_BYTES:
-        raise ValueError("File too large. Please upload an image ≤ 8 MB.")
-    with Image.open(path) as im:
-        w, h = im.size
-        if w > MAX_SIDE or h > MAX_SIDE:
-            raise ValueError("Image too large (dimensions). Please keep ≤ 4096×4096.")
-    return path
-
-# =========================
-# Inference (with optional simple TTA)
-# =========================
 def infer(image_file, threshold=0.5, tta=False):
     try:
-        if PREDICTOR is None:
-            raise RuntimeError("Model failed to load. Check model artifacts and environment.")
-
-        # 1) Validate & open
-        path = validate_image(image_file)
+        path = image_file.name if hasattr(image_file, "name") else str(image_file)
         orig = Image.open(path).convert("RGB")
         vis = preprocess_for_viz(orig.copy())
 
-        # 2) Build test dataframe for predictor
         df = pd.DataFrame([{"image": path}])
-
-        # 3) Predict proba for lipstick=1 (binary)
         proba_main = PREDICTOR.predict_proba(df)
         row = proba_main.iloc[0]
+
         p1 = None
         for k in [1, "1", "lipstick", "Lipstick", "positive", "True"]:
             if k in row.index:
-                p1 = float(row[k]); break
+                p1 = float(row[k])
+                break
         if p1 is None:
             if len(row.index) == 2:
                 p1 = float(max(row.values))
             else:
                 p1 = float(row.values[0])
 
-        # 4) Optional simple TTA: average with horizontally flipped image prediction
         if tta:
             flipped = orig.transpose(Image.FLIP_LEFT_RIGHT)
             flip_tmp = pathlib.Path(path).with_suffix(".flip_tmp.jpg")
@@ -161,14 +95,15 @@ def infer(image_file, threshold=0.5, tta=False):
             p1_flip = None
             for k in [1, "1", "lipstick", "Lipstick", "positive", "True"]:
                 if k in row_flip.index:
-                    p1_flip = float(row_flip[k]); break
+                    p1_flip = float(row_flip[k])
+                    break
             if p1_flip is None:
-                p1_flip = float(max(row_flip.values)) if len(row_flip.index)==2 else float(row_flip.values[0])
+                if len(row_flip.index) == 2:
+                    p1_flip = float(max(row_flip.values))
+                else:
+                    p1_flip = float(row_flip.values[0])
             p1 = float((p1 + p1_flip) / 2.0)
-            try:
-                os.remove(flip_tmp)
-            except Exception:
-                pass
+            os.remove(flip_tmp)
 
         p1 = float(np.clip(p1, 0.0, 1.0))
         p0 = float(1.0 - p1)
@@ -180,9 +115,6 @@ def infer(image_file, threshold=0.5, tta=False):
         tb = traceback.format_exc(limit=1)
         return None, None, {"error": f"{type(e).__name__}: {e}"}, f"Failed: {type(e).__name__}: {e}\n{tb}"
 
-# =========================
-# Build examples from dataset (if available)
-# =========================
 def get_examples(n=3):
     try:
         ds = load_dataset(DATASET_REPO, split="validation")
@@ -211,12 +143,13 @@ def get_examples(n=3):
 
 EXAMPLES = get_examples(3)
 
-# =========================
-# Gradio UI
-# =========================
 with gr.Blocks() as demo:
-    gr.Markdown(f"# {TITLE}")
-    gr.Markdown(DESC)
+    gr.Markdown("# 💄 Lipstick Detection (AutoGluon — EfficientNet-B0)")
+    gr.Markdown(
+        "- Upload a face image; the model predicts **Lipstick** vs **No Lipstick**.  \n"
+        "- Left: original; Right: the **preprocessed** 224×224 view seen by the model.  \n"
+        "- This is a **teaching demo**; don’t use for real decisions."
+    )
 
     with gr.Row():
         with gr.Column():
@@ -239,6 +172,4 @@ with gr.Blocks() as demo:
         cache_examples=False
     )
 
-# In Spaces, share=True is not required; leaving default.
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()