classifier / app.py
drixo's picture
Update app.py
4d642ec verified
import gradio as gr
from duckduckgo_search import DDGS
from fastcore.all import *
from fastdownload import download_url
from fastai.vision.all import *
from PIL import Image
from pathlib import Path
# Define search function using DuckDuckGo
ddgs = DDGS()
def search_images(term, max_images=30):
print(f"Searching for '{term}'")
return L(ddgs.images(term, max_results=max_images)).itemgot('image')
# Create a folder for storing images
path = Path('images')
path.mkdir(exist_ok=True)
# Download example image: beaver
beaver_url = search_images('beaver photo', max_images=1)[0]
download_url(beaver_url, path/'beaver.jpg', show_progress=False)
# Download another example image: platypus
platypus_url = search_images('platypus photo', max_images=1)[0]
download_url(platypus_url, path/'platypus.jpg', show_progress=False)
# Show a thumbnail of the platypus image
Image.open(path/'platypus.jpg').thumbnail((256,256))
# Remove any corrupt images
failed = verify_images(get_image_files(path))
failed.map(Path.unlink)
# Prepare DataLoaders (make sure `images/` has subfolders of labeled images)
dls = DataBlock(
blocks=(ImageBlock, CategoryBlock),
get_items=get_image_files,
splitter=RandomSplitter(valid_pct=0.2, seed=42),
get_y=parent_label,
item_tfms=[Resize(192, method='squish')]
).dataloaders(path, bs=32)
# Show a sample batch
dls.show_batch(max_n=6)
# Train a model
learn = vision_learner(dls, resnet18, metrics=error_rate)
learn.fine_tune(3)
# Prediction function
def predict_species(img):
is_sheep, _, probs = learn.predict(img)
return f"This looks like a: {is_sheep}. Probability it's a beaver: {probs[0]:.4f}"
# Define Gradio interface
demo = gr.Interface(fn=predict_species, inputs=gr.Image(type="pil"), outputs="text")
demo.launch()