Spaces:
Sleeping
Sleeping
| import io, os, tempfile, pathlib, zipfile, shutil | |
| from PIL import Image, ImageOps | |
| import pandas as pd # needed to hold predictor output, but we won't return any DataFrame | |
| from huggingface_hub import hf_hub_download | |
| import autogluon.multimodal as ag | |
| import gradio as gr | |
| MODEL_REPO_ID = "its-zion-18/sign-image-autogluon-predictor" | |
| ZIP_FILENAME = "autogluon_image_predictor_dir.zip" | |
| CACHE_DIR = pathlib.Path("hf_assets") | |
| EXTRACT_DIR = CACHE_DIR / "predictor_native" | |
| CLASS_LABELS = {0: "Other sign", 1: "Stop sign"} | |
| DEFAULT_SIZE = 224 | |
| def _prepare_predictor_dir() -> str: | |
| CACHE_DIR.mkdir(parents=True, exist_ok=True) | |
| local_zip = hf_hub_download( | |
| repo_id=MODEL_REPO_ID, | |
| filename=ZIP_FILENAME, | |
| repo_type="model", | |
| 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()) | |
| return str(contents[0] if len(contents) == 1 and contents[0].is_dir() else EXTRACT_DIR) | |
| PREDICTOR = ag.MultiModalPredictor.load(_prepare_predictor_dir()) | |
| def preprocess_for_model(img: Image.Image, size: int = DEFAULT_SIZE) -> Image.Image: | |
| return ImageOps.fit(img, (size, size), method=Image.BICUBIC, centering=(0.5, 0.5)) | |
| # ------------------------- | |
| # Prediction (return 4 values to match your outputs) | |
| # ------------------------- | |
| def predict(img: Image.Image): | |
| if img is None: | |
| return None, None, {"Error": 1.0}, "**Error:** No image provided." | |
| pre_img = preprocess_for_model(img, DEFAULT_SIZE) | |
| # temp save for AutoGluon | |
| tmpdir = pathlib.Path(tempfile.mkdtemp()) | |
| img_path = tmpdir / "input.png" | |
| pre_img.save(img_path) | |
| df = pd.DataFrame({"image": [str(img_path)]}) | |
| proba_df = PREDICTOR.predict_proba(df) | |
| row = proba_df.iloc[0] | |
| # Handle int or str column names | |
| p0 = float(row.get(0, row.get("0", 0.0))) | |
| p1 = float(row.get(1, row.get("1", 0.0))) | |
| pretty = {CLASS_LABELS[0]: p0, CLASS_LABELS[1]: p1} | |
| status = "" # empty if OK | |
| # EXACTLY 4 returns to match your outputs list | |
| return img, pretty, status | |
| EXAMPLES = [ | |
| "https://drive.google.com/uc?id=1OwjdQdiclDOCuUirVsCX-_J1q1Y_G-Xs", | |
| "https://drive.google.com/uc?id=1g9XwRJGcqYwJ_mFtHj25ZE1w2VZPkBiR", | |
| "https://drive.google.com/uc?id=1MH5aIvHWKdamwuzdQPWlNwaoRE3c2bVb" | |
| ] | |
| demo = gr.Interface( | |
| fn=predict, | |
| inputs=gr.Image(type="pil", label="Input image (PNG/JPG)", sources=["upload", "webcam"]), | |
| outputs=[ | |
| gr.Image(label="Preprocessed"), | |
| gr.Label(num_top_classes=2, label="Class probabilities"), | |
| gr.Markdown(label="Status"), | |
| ], | |
| title="🪧 Sign Identification", | |
| description="Upload a PNG/JPG. Predicts Stop vs Other sign.", | |
| allow_flagging="never", | |
| examples=EXAMPLES, | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |