Spaces:
Sleeping
Sleeping
File size: 4,767 Bytes
ecb2e52 e16f7b0 ecb2e52 5fb09aa ecb2e52 |
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 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
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
# --- Model Loading ---
def _prepare_predictor_dir() -> str:
"""Downloads and extracts the model files from Hugging Face."""
# Add a check to clear the cache before download
if CACHE_DIR.exists():
print(f"Clearing cache directory: {CACHE_DIR}")
shutil.rmtree(CACHE_DIR)
CACHE_DIR.mkdir(parents=True, exist_ok=True)
# Hardcoded Hub model (native zip)
MODEL_REPO_ID = "apsora/autoML_images_data"
ZIP_FILENAME = "autogluon_image_predictor_dir.zip"
HF_TOKEN = os.getenv("HF_TOKEN", None)
# Local cache/extract dirs
CACHE_DIR = pathlib.Path("hf_assets")
EXTRACT_DIR = CACHE_DIR / "predictor_native"
# Download & load the native predictor
def _prepare_predictor_dir() -> str:
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)
# Explicit class labels (edit copy as desired)
CLASS_LABELS = {0: "π« Not a tomato", 1: "π
Tomato"}
# Helper to map model class -> human label
def _human_label(c):
try:
ci = int(c)
return CLASS_LABELS.get(ci, str(c))
except Exception:
return CLASS_LABELS.get(c, str(c))
# Do the prediction!
def do_predict(pil_img: PIL.Image.Image):
# Make sure there's actually an image to work with
if pil_img is None:
return "No image provided.", {}, pandas.DataFrame(columns=["Predicted label", "Confidence (%)"])
# IF we have something to work with, save it and prepare the input
tmpdir = pathlib.Path(tempfile.mkdtemp())
img_path = tmpdir / "input.png"
pil_img.save(img_path)
df = pandas.DataFrame({"image": [str(img_path)]}) # For AutoGluon expected input format
# For class probabilities
proba_df = PREDICTOR.predict_proba(df)
# For user-friendly column names
proba_df = proba_df.rename(columns={0: "π« Not a tomato (0)", 1: "π
Tomato (1)"})
row = proba_df.iloc[0]
# For pretty ranked dict expected by gr.Label
pretty_dict = {
"π« Not a tomato": float(row.get("π« Not a tomato (0)", 0.0)),
"π
Tomato": float(row.get("π
Tomato (1)", 0.0)),
}
return pretty_dict
# Representative example images! These can be local or links.
EXAMPLES = [
["https://dengarden.com/.image/w_1920,q_auto:good,c_limit/MTk3NDQ3MTk3NDE4MDcxMDQ2/how-to-get-the-highest-yield-and-best-flavor-from-tomatoes.jpg"],
["https://www.greenlanedelivery.com/cdn/shop/products/Grapes_White_SL_1200x1200.jpg?v=1671549475"],
["https://agrinigeriaprodsa.blob.core.windows.net/agrifarmer/a8738a87-3e02-4d1c-8ba7-e028205ee6bb.jpg"]
]
# Gradio UI
with gradio.Blocks() as demo:
# Provide an introduction
gradio.Markdown("# Tomato or No Tomato?")
gradio.Markdown("""
This is a simple app that demonstrates how to use an autogluon multimodal
predictor in a gradio space to predict the contents of a picture. To use,
just upload a photo. The result should be generated automatically.
""")
# Interface for the incoming image
image_in = gradio.Image(type="pil", label="Input image", sources=["upload", "webcam"])
# Interface elements to show htte result and probabilities
proba_pretty = gradio.Label(num_top_classes=2, label="Class probabilities")
# Whenever a new image is uploaded, update the result
image_in.change(fn=do_predict, inputs=[image_in], outputs=[proba_pretty])
# For clickable example images
gradio.Examples(
examples=EXAMPLES,
inputs=[image_in],
label="Representative examples",
examples_per_page=8,
cache_examples=False,
)
if __name__ == "__main__":
demo.launch() |