File size: 3,969 Bytes
aabac73
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import os
import shutil
import zipfile
import pathlib
import tempfile

import gradio
import pandas
import PIL.Image

import huggingface_hub
import autogluon.multimodal

# Hardcoded Hub model (native zip)
MODEL_REPO_ID = "aedupuga/image-autogluon-predictor"
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: "Fiction", 1: "Nonfiction"}

# 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: "Fiction", 1: "Nonfiction"})
    row = proba_df.iloc[0]

    # For pretty ranked dict expected by gr.Label
    pretty_dict = {
        "Fiction": float(row.get("Fiction", 0.0)),
        "Nonfiction": float(row.get("Nonfiction", 0.0)),
    }

    return pretty_dict

# Representative example images! These can be local or links.
EXAMPLES = [
    ["https://pictures.abebooks.com/inventory/31174850639.jpg"],
    ["https://res.cloudinary.com/bloomsbury-atlas/image/upload/w_568,c_scale,dpr_1.5/jackets/9781408872925.jpg"],
    ["https://images3.penguinrandomhouse.com/cover/9780593191897"]
]

# Gradio UI
with gradio.Blocks() as demo:

    # Provide an introduction
    gradio.Markdown("# Judge a Book by its Cover: Fiction or Nonfiction?")
    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()