Upload app.py with huggingface_hub

app.py

@@ -1,18 +1,16 @@
-import pathlib, shutil, zipfile
+import pathlib, shutil, zipfile, tempfile
 import pandas
 import gradio
 import huggingface_hub
-import autogluon.tabular
+import autogluon.multimodal
+import PIL.Image
 
-MODEL_REPO_ID = "jennifee/classical_automl_model"
-ZIP_FILENAME  = "autogluon_predictor_dir.zip"
+MODEL_REPO_ID = "george2cool36/hw2_image_automl_autogluon"
+ZIP_FILENAME  = "ag_image_predictor_dir.zip"
 CACHE_DIR = pathlib.Path("hf_assets")
 EXTRACT_DIR = CACHE_DIR / "predictor_native"
 
-FEATURE_COLS_MODEL = ['phone_hours', 'computer_hours', 'device_count', 'sleep_quality', 'sleep_time', 'sleep_hours']
-TARGET_COL = "use_before_bed"
-SLEEP_QUALITY_LABELS = ['good', 'medium', 'bad']
-USE_BEFORE_BED_LABELS = {0: 'No', 1: 'Yes'}
+CLASS_LABELS = {0: "🛑 Has Stop Sign", 1: "✅ No Stop Sign"}
 
 def _prepare_predictor_dir() -> str:
     CACHE_DIR.mkdir(parents=True, exist_ok=True)
@@ -33,86 +31,66 @@ def _prepare_predictor_dir() -> str:
     return str(predictor_root)
 
 PREDICTOR_DIR = _prepare_predictor_dir()
-PREDICTOR = autogluon.tabular.TabularPredictor.load(PREDICTOR_DIR, require_py_version_match=False)
+PREDICTOR = autogluon.multimodal.MultiModalPredictor.load(PREDICTOR_DIR)
 
 def _human_label(c):
     try:
         ci = int(c)
-        if ci in USE_BEFORE_BED_LABELS:
-            return USE_BEFORE_BED_LABELS[ci]
+        return CLASS_LABELS.get(ci, str(c))
     except Exception:
-        pass
-    return str(c)
-
-def do_predict(phone_hours, computer_hours, device_count, sleep_quality_label, sleep_time, sleep_hours):
-    row = {
-        "phone_hours": float(phone_hours),
-        "computer_hours": float(computer_hours),
-        "device_count": int(device_count),
-        "sleep_quality": sleep_quality_label,
-        "sleep_time": int(sleep_time),
-        "sleep_hours": float(sleep_hours),
-    }
-    X = pandas.DataFrame([row], columns=FEATURE_COLS_MODEL)
+        return CLASS_LABELS.get(c, str(c))
 
-    pred_series = PREDICTOR.predict(X)
-    raw_pred = pred_series.iloc[0]
-    pred_label = _human_label(raw_pred)
+def do_predict(pil_img: PIL.Image.Image):
+    if pil_img is None:
+        return {}, "No image provided."
 
-    proba = None
-    try:
-        proba = PREDICTOR.predict_proba(X)
-        if isinstance(proba, pandas.Series):
-            proba = proba.to_frame().T
-    except Exception as e:
-        print(f"Error calculating probabilities: {e}")
-        proba = None
-
-    proba_dict = None
-    if proba is not None:
-        row0 = proba.iloc[0]
-        tmp = {}
-        for cls in [0, 1]:
-            val = None
-            if cls in row0.index:
-                val = row0[cls]
-            elif str(cls) in row0.index:
-                val = row0[str(cls)]
-            if val is not None:
-                key = _human_label(cls)
-                tmp[key] = float(tmp.get(key, 0.0)) + float(val)
-        if tmp:
-            proba_dict = dict(sorted(tmp.items(), key=lambda kv: kv[1], reverse=True))
+    tmpdir = pathlib.Path(tempfile.mkdtemp())
+    img_path = tmpdir / "input.png"
+    pil_img.save(img_path)
+
+    df = pandas.DataFrame({"image": [str(img_path)]})
+
+    proba_df = PREDICTOR.predict_proba(df)
+    proba_df = proba_df.rename(columns={0: "🛑 Has Stop Sign (0)", 1: "✅ No Stop Sign (1)"})
+    row = proba_df.iloc[0]
+
+    pretty_dict = {
+        "🛑 Has Stop Sign": float(row.get("🛑 Has Stop Sign (0)", 0.0)),
+        "✅ No Stop Sign": float(row.get("✅ No Stop Sign (1)", 0.0)),
+    }
+
+    predicted_class = PREDICTOR.predict(df).iloc[0]
+    pred_label = _human_label(predicted_class)
 
     md = f"**Prediction:** {pred_label}"
-    if proba_dict:
-        md += f"  \n**Confidence:** {round(proba_dict.get(pred_label, 0.0) * 100, 2)}%"
+    if pretty_dict:
+        md += f"  \n**Confidence:** {round(pretty_dict.get(pred_label, 0.0) * 100, 2)}%"
+
+
+    return pretty_dict, md
 
-    return proba_dict, md
+EXAMPLES = [
+    ["https://www.kingsrivercasting.com/images/stories/virtuemart/product/STOP%20SIGN%20(5).jpg"],
+    ["https://www.trafficsafetywarehouse.com/Resources/images/traffic-sign-shapes.jpeg"],
+    ["https://di-uploads-pod16.dealerinspire.com/toyotaofnorthcharlotte/uploads/2020/08/yield-road-sign.jpg"]
+]
 
-EXAMPLES = [[3.5, 5.0, 3, 'good', 23, 7.0], [4.2, 6.5, 3, 'medium', 0, 6.5], [5.0, 4.0, 4, 'bad', 1, 6.0], [2.0, 7.5, 3, 'good', 22, 7.5], [3.8, 6.0, 3, 'medium', 0, 6.0]]
 
 with gradio.Blocks() as demo:
-    gradio.Markdown("# Predict Phone Use Before Bed")
+    gradio.Markdown("# Has Stop Sign or Not?")
     gradio.Markdown(
-        "This app predicts whether a student uses their phone before bed based on their sleeping habits."
-        "\nEnter the student's sleeping habits below to get a prediction."
+        "This is a simple app that demonstrates how to use an autogluon multimodal"
+        "predictor in a gradio space to predict whether an image contains a stop sign. To use,"
+        "just upload a photo. The result should be generated automatically."
     )
 
-    with gradio.Row():
-        phone_hours = gradio.Number(value=3.5, precision=1, label=FEATURE_COLS_MODEL[0])
-        computer_hours = gradio.Number(value=5.0, precision=1, label=FEATURE_COLS_MODEL[1])
-        device_count = gradio.Number(value=3, precision=0, label=FEATURE_COLS_MODEL[2])
+    image_in = gradio.Image(type="pil", label="Input image", sources=["upload", "webcam"])
 
-    with gradio.Row():
-        sleep_quality_label = gradio.Radio(choices=SLEEP_QUALITY_LABELS, value="good", label=FEATURE_COLS_MODEL[3])
-        sleep_time = gradio.Number(value=23, precision=0, label=FEATURE_COLS_MODEL[4])
-        sleep_hours = gradio.Number(value=7.0, precision=1, label=FEATURE_COLS_MODEL[5])
-
-    proba_pretty = gradio.Label(num_top_classes=2, label="Probability of Using Phone Before Bed")
+    proba_pretty = gradio.Label(num_top_classes=2, label="Class probabilities")
     prediction_output = gradio.Markdown()
 
-    inputs = [phone_hours, computer_hours, device_count, sleep_quality_label, sleep_time, sleep_hours]
+
+    inputs = [image_in]
     outputs = [proba_pretty, prediction_output]
     for comp in inputs:
         comp.change(fn=do_predict, inputs=inputs, outputs=outputs)
@@ -121,7 +99,7 @@ with gradio.Blocks() as demo:
         examples=EXAMPLES,
         inputs=inputs,
         label="Representative examples",
-        examples_per_page=5,
+        examples_per_page=8,
         cache_examples=False,
     )