File size: 3,515 Bytes
9875b71 |
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 |
# !pip install -q autogluon.tabular gradio huggingface_hub pandas
from huggingface_hub import hf_hub_download
from autogluon.tabular import TabularPredictor
import pathlib, shutil, zipfile
import pandas as pd
import gradio as gr
# ---------------- 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"
)
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)
PREDICTOR = load_predictor()
# ---------------- 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"])
row["Studs"] = int(round(float(row["Studs"]))) # gr.Number returns float
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:
X = _cast_and_rename({
"Length": length, "Height": height, "Width": width, "Studs": studs
})
pred = PREDICTOR.predict(X)
pred_val = pred.iloc[0] if hasattr(pred, "iloc") else pred
# 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:
return {"prediction": str(pred_val)}
except Exception as e:
import traceback
return {"error": f"{type(e).__name__}: {e}", "traceback": traceback.format_exc()}
# ---------------- Quick test (uses correct names) ----------------
# test_X = _cast_and_rename({"Length": 4, "Height": 1.2, "Width": 2, "Studs": 4})
# print("Prediction:", PREDICTOR.predict(test_X))
# print("Probabilities:\n", PREDICTOR.predict_proba(test_X))
# ---------------- Gradio (Interface version) ----------------
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="🧱 LEGO Brick Classifier",
description="Predicts whether a LEGO piece is Standard, Flat, or Sloped."
)
demo.launch()
|