File size: 1,390 Bytes
e88ba5b ac5e28e e88ba5b ac5e28e adfe3bd e88ba5b adfe3bd e88ba5b adfe3bd e88ba5b adfe3bd e88ba5b adfe3bd e88ba5b | 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 | # app.py
import os, shutil
# keep all caches ephemeral
os.environ["HF_HOME"]="/tmp/hf"
os.environ["TRANSFORMERS_CACHE"]="/tmp/hf/transformers"
os.environ["TORCH_HOME"]="/tmp/torch"
os.environ["PIP_NO_CACHE_DIR"]="1"
for d in ["/home/user/.cache", "/home/user/.fastai", "/home/user/.torch"]:
shutil.rmtree(d, ignore_errors=True)
from functools import lru_cache
from huggingface_hub import hf_hub_download
from fastai.vision.all import load_learner
import gradio as gr
REPO_ID = "daleef/my-fastai-bear" # your model repo
FILENAME = "model.pkl"
@lru_cache(maxsize=1)
def get_learner():
pkl_path = hf_hub_download(
repo_id=REPO_ID,
filename=FILENAME,
local_dir="/tmp/model",
local_dir_use_symlinks=False
)
return load_learner(pkl_path)
def classify_image(img):
learn = get_learner()
pred, idx, probs = learn.predict(img)
classes = learn.dls.vocab
return {c: float(probs[i]) for i, c in enumerate(classes)}
demo = gr.Interface(
fn=classify_image,
inputs=gr.Image(type="pil", height=192, width=192, label="Upload an image"),
outputs=gr.Label(num_top_classes=3),
title="Fastai Bear Classifier",
description="Upload a bear image to classify."
# Tip: remove examples unless those files exist in the repo
# examples=["black.jpg","bear.jpg","teddy.jpg"]
)
if __name__ == "__main__":
demo.launch() |