File size: 2,196 Bytes
9482d82
 
 
 
 
 
 
 
8cdc0da
9482d82
 
 
 
 
 
 
a2c1d64
9482d82
 
 
 
 
 
 
8cdc0da
9482d82
 
 
a2c1d64
9482d82
847ccdc
9482d82
 
 
 
b54e1b9
9482d82
 
 
 
a7249a4
9482d82
 
 
 
 
 
 
 
6a7b3f5
9482d82
 
8cdc0da
97f4fd9
8cdc0da
 
 
9482d82
 
 
 
 
852a7c0
9482d82
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
import gradio
import pandas as pd
import PIL.Image
import tempfile, pathlib, requests, io

from huggingface_hub import snapshot_download
from autogluon.multimodal import MultiModalPredictor

# Hugging Face model repo
MODEL_REPO_ID = "scottymcgee/image-classifier"

# Download full repo snapshot locally
repo_dir = snapshot_download(repo_id=MODEL_REPO_ID)

# Load predictor directly
PREDICTOR = MultiModalPredictor.load(repo_dir)
TMP_PATH = pathlib.Path("tmp_input.png")

# Explicit class labels
CLASS_LABELS = {0: "🚫 No Stop Sign", 1: "🛑 Stop Sign"}

def do_predict(pil_img: PIL.Image.Image):
    if pil_img is None:
        return {"Error": 1.0}
        
    pil_resized = pil_img.convert("RGB").resize((256, 256))
    
    # Save to tmp for AutoGluon input
    pil_resized.save(TMP_PATH)

    df = pd.DataFrame({"image": [str(TMP_PATH)]})

    # Predict probabilities
    proba_df = PREDICTOR.predict_proba(df)

    # Map columns to names
    pretty_dict = {}
    for idx, col in enumerate(proba_df.columns):
        label = CLASS_LABELS.get(idx, str(col))
        pretty_dict[label] = float(proba_df.iloc[0][col])
    return pil_resized, pretty_dict

EXAMPLES = [
    ["examples/stop1.jpg"],
    ["examples/no_stop1.jpg"],
    ["examples/stop2.jpg"],
]

with gradio.Blocks() as demo:
    gradio.Markdown("# Stop Sign Classifier")
    gradio.Markdown("Upload a traffic image. The model shows the **original**, the **preprocessed**, and the prediction.")

    with gradio.Row():
        threshold = gradio.Slider(
            0.0, 1.0, value=0.2, step=0.05, label="Confidence threshold"
        )

    with gradio.Row():
        image_in = gradio.Image(type="pil", label="Upload image", sources=["upload", "webcam"])
        proc_out = gradio.Image(type="pil", label="Preprocessed (256x256)")
        proba_pretty = gradio.Label(num_top_classes=2, label="Class probabilities")

    image_in.change(fn=do_predict, inputs=[image_in], outputs=[proc_out, proba_pretty])

    gradio.Examples(
        examples=EXAMPLES,
        inputs=[image_in],
        outputs=[proc_out, proba_pretty],
        fn=do_predict,
        cache_examples=False,
    )

if __name__ == "__main__":
    demo.launch()