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()