Spaces:
Sleeping
Sleeping
File size: 3,722 Bytes
7084745 0b09d2c 7084745 0b09d2c 7084745 0b09d2c 7084745 0b09d2c 7084745 0b09d2c 7084745 0b09d2c 7084745 0b09d2c 7084745 0b09d2c 7084745 0b09d2c 7084745 0b09d2c 7084745 0b09d2c 7084745 0b09d2c 7084745 0b09d2c 7084745 0b09d2c |
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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
import pathlib, shutil, zipfile, os, traceback
import pandas as pd
import gradio as gr
from huggingface_hub import hf_hub_download
from autogluon.tabular import TabularPredictor
# ---------------- UI copy ----------------
TITLE = "🧱 LEGO Brick Classifier"
DESC = "Predicts whether a LEGO piece is Standard, Flat, or Sloped from basic dimensions."
# ---------------- Settings ----------------
MODEL_REPO_ID = "Iris314/classical-automl-model"
ZIP_FILENAME = "lego_predictor_dir.zip"
# UI → model feature name mapping
COLUMN_ALIAS = {
"Length": "Max Length (cm)",
"Height": "Max Height (cm)",
"Width": "Width (cm)",
"Studs": "Studs",
}
FEATURE_COLS_UI = ["Length", "Height", "Width", "Studs"]
# ---------------- Load predictor ----------------
CACHE_DIR = pathlib.Path("hf_cache"); EXTRACT_DIR = CACHE_DIR / "predictor"
CACHE_DIR.mkdir(exist_ok=True, parents=True)
def load_predictor():
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)
with zipfile.ZipFile(local_zip, "r") as zf:
zf.extractall(EXTRACT_DIR)
kids = list(EXTRACT_DIR.iterdir())
path = kids[0] if len(kids) == 1 and kids[0].is_dir() else EXTRACT_DIR
return TabularPredictor.load(str(path), require_py_version_match=False)
try:
PREDICTOR = load_predictor()
except Exception as e:
PREDICTOR = None
print("Failed to load predictor:", e)
# ---------------- Helpers ----------------
def _cast_and_rename(row_dict):
row = dict(row_dict)
row["Length"] = float(row["Length"])
row["Height"] = float(row["Height"])
row["Width"] = float(row["Width"])
# gr.Number returns float; round & cast for integer feature
row["Studs"] = int(round(float(row["Studs"])))
X_ui = pd.DataFrame([row], columns=FEATURE_COLS_UI)
X_model = X_ui.rename(columns=COLUMN_ALIAS)
return X_model
def classify_brick(length, height, width, studs):
try:
if PREDICTOR is None:
raise RuntimeError("Model failed to load on startup. Check model artifact path & runtime deps.")
X = _cast_and_rename({
"Length": length, "Height": height, "Width": width, "Studs": studs
})
# Try probabilities; fall back to label
try:
proba = PREDICTOR.predict_proba(X)
s = proba.iloc[0] if hasattr(proba, "iloc") else proba
s = s.sort_values(ascending=False)
s.index = [str(k) for k in s.index] # ensure JSON-serializable keys
return {k: float(v) for k, v in s.items()}
except Exception:
pred = PREDICTOR.predict(X)
pred_val = pred.iloc[0] if hasattr(pred, "iloc") else pred
return {"prediction": str(pred_val)}
except Exception as e:
return {
"error": f"{type(e).__name__}: {e}",
"traceback": traceback.format_exc(limit=1)
}
# ---------------- Gradio ----------------
demo = gr.Interface(
fn=classify_brick,
inputs=[
gr.Slider(1, 10, step=0.1, value=4, label="Length"),
gr.Slider(0.2, 5, step=0.1, value=1.2, label="Height"),
gr.Slider(1, 10, step=0.1, value=2, label="Width"),
gr.Number(value=4, precision=0, label="Studs"),
],
outputs=gr.Label(num_top_classes=3, label="Predicted Class / Probabilities"),
examples=[[4, 1.2, 2, 4], [2, 0.6, 2, 2], [3, 2.0, 2, 2]],
title=TITLE,
description=DESC
)
if __name__ == "__main__":
# In Spaces, no share=True needed
demo.launch()
|