Spaces:
Runtime error
Runtime error
| # ================================================================ | |
| # HW3: Driving a Stop Sign Image Classifier with Gradio | |
| # | |
| # Author: Your Name | |
| # Course: 24679 - Designing and Deploying AI/ML Systems | |
| # Dataset: Binary stop sign dataset (class_1 = stop sign, class_0 = not stop sign) | |
| # Task: Image classification deployed via Hugging Face Space | |
| # | |
| # Acknowledgments: | |
| # - Model trained by a classmate in Homework 2 | |
| # - Deployment scaffold and documentation supported with AI assistance (ChatGPT, OpenAI) | |
| # - Reference: Class-provided notebook "image gradio.ipynb" | |
| # ================================================================ | |
| import os # For reading environment variables | |
| import shutil # For directory cleanup | |
| import zipfile # For extracting model archives | |
| import pathlib # For path manipulations | |
| import tempfile # For creating temporary files/directories | |
| import gradio # For interactive UI | |
| import pandas # For tabular data handling | |
| import PIL.Image # For image I/O | |
| import huggingface_hub # For downloading model assets | |
| import autogluon.multimodal # For loading AutoGluon image classifier | |
| # ------------------------- | |
| # Hugging Face model setup | |
| # ------------------------- | |
| MODEL_REPO_ID = "cassieli226/sign-identification-automl" # <- update with teammate’s repo | |
| ZIP_FILENAME = "autogluon_predictor_dir.zip" | |
| HF_TOKEN = os.getenv("HF_TOKEN", None) | |
| CACHE_DIR = pathlib.Path("hf_assets") | |
| EXTRACT_DIR = CACHE_DIR / "predictor_native" | |
| def _prepare_predictor_dir() -> str: | |
| """Download and extract the AutoGluon predictor directory from Hugging Face.""" | |
| CACHE_DIR.mkdir(parents=True, exist_ok=True) | |
| local_zip = huggingface_hub.hf_hub_download( | |
| repo_id=MODEL_REPO_ID, | |
| filename=ZIP_FILENAME, | |
| repo_type="model", | |
| token=HF_TOKEN, | |
| local_dir=str(CACHE_DIR), | |
| local_dir_use_symlinks=False, | |
| ) | |
| if EXTRACT_DIR.exists(): | |
| shutil.rmtree(EXTRACT_DIR) | |
| EXTRACT_DIR.mkdir(parents=True, exist_ok=True) | |
| with zipfile.ZipFile(local_zip, "r") as zf: | |
| zf.extractall(str(EXTRACT_DIR)) | |
| contents = list(EXTRACT_DIR.iterdir()) | |
| predictor_root = contents[0] if (len(contents) == 1 and contents[0].is_dir()) else EXTRACT_DIR | |
| return str(predictor_root) | |
| PREDICTOR_DIR = _prepare_predictor_dir() | |
| PREDICTOR = autogluon.multimodal.MultiModalPredictor.load(PREDICTOR_DIR) | |
| # ------------------------- | |
| # Class labels | |
| # ------------------------- | |
| CLASS_LABELS = { | |
| 0: "🚦 Not a Stop Sign", | |
| 1: "🛑 Stop Sign" | |
| } | |
| def _human_label(c): | |
| try: | |
| ci = int(c) | |
| return CLASS_LABELS.get(ci, str(c)) | |
| except Exception: | |
| return CLASS_LABELS.get(c, str(c)) | |
| # ------------------------- | |
| # Prediction + preprocessing | |
| # ------------------------- | |
| def do_predict(pil_img: PIL.Image.Image): | |
| """Run prediction on an uploaded image. Returns original, preprocessed, and probabilities.""" | |
| if pil_img is None: | |
| return None, None, "No image provided.", {} | |
| try: | |
| # --- Save to temp path for AutoGluon --- | |
| tmpdir = pathlib.Path(tempfile.mkdtemp()) | |
| img_path = tmpdir / "input.png" | |
| pil_img.save(img_path) | |
| # --- Preprocess (resize to 224x224 for visualization only) --- | |
| preprocessed_img = pil_img.copy() | |
| preprocessed_img = preprocessed_img.resize((224, 224)) | |
| # --- Build input dataframe for AutoGluon --- | |
| df = pandas.DataFrame({"image": [str(img_path)]}) | |
| # --- Predict probabilities --- | |
| proba_df = PREDICTOR.predict_proba(df) | |
| # Rename for clarity | |
| proba_df = proba_df.rename(columns={ | |
| 0: "🚦 Not a Stop Sign (0)", | |
| 1: "🛑 Stop Sign (1)" | |
| }) | |
| row = proba_df.iloc[0] | |
| pretty_dict = { | |
| "🚦 Not a Stop Sign": float(row.get("🚦 Not a Stop Sign (0)", 0.0)), | |
| "🛑 Stop Sign": float(row.get("🛑 Stop Sign (1)", 0.0)), | |
| } | |
| return pil_img, preprocessed_img, "Prediction complete", pretty_dict | |
| except Exception as e: | |
| return None, None, f"Error: {str(e)}", {} | |
| # ------------------------- | |
| # Example images | |
| # ------------------------- | |
| EXAMPLES = [ | |
| ["https://upload.wikimedia.org/wikipedia/commons/thumb/f/f9/STOP_sign.jpg/640px-STOP_sign.jpg"], | |
| ["https://upload.wikimedia.org/wikipedia/commons/1/19/Swiss_Frutiger_Traffic_Sign.jpg"] | |
| ] | |
| # ------------------------- | |
| # Gradio interface | |
| # ------------------------- | |
| with gradio.Blocks() as demo: | |
| gradio.Markdown("# 🛑 Stop Sign Detector") | |
| gradio.Markdown(""" | |
| Upload a road scene or traffic sign image, and this app will classify whether | |
| a stop sign is present. The interface shows both the original image and the | |
| preprocessed image (224x224) that the model actually sees. | |
| """) | |
| with gradio.Row(): | |
| image_in = gradio.Image( | |
| type="pil", | |
| label="Upload or capture an image", | |
| sources=["upload", "webcam"], | |
| image_mode="RGB" | |
| ) | |
| with gradio.Row(): | |
| orig_out = gradio.Image(type="pil", label="Original Image") | |
| preproc_out = gradio.Image(type="pil", label="Preprocessed Image (224x224)") | |
| status_out = gradio.Textbox(label="Status") | |
| proba_pretty = gradio.Label(num_top_classes=2, label="Class probabilities") | |
| image_in.change( | |
| fn=do_predict, | |
| inputs=[image_in], | |
| outputs=[orig_out, preproc_out, status_out, proba_pretty] | |
| ) | |
| gradio.Examples( | |
| examples=EXAMPLES, | |
| inputs=[image_in], | |
| label="Representative examples", | |
| examples_per_page=3, | |
| cache_examples=False, | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |