|
from duckduckgo_search import ddg_images |
|
from fastai.vision.all import download_images, resize_images, verify_images, get_image_files, ImageBlock, \ |
|
CategoryBlock, RandomSplitter, parent_label, ResizeMethod, Resize, vision_learner, resnet18, error_rate, \ |
|
L, Path, DataBlock |
|
|
|
|
|
def search_images(search_term, max_images=30): |
|
print(f"Searching for '{search_term}'") |
|
return L(ddg_images(search_term, max_results=max_images)).itemgot('image') |
|
|
|
|
|
def search_and_populate(search_term, category, file_path, max_images=30): |
|
dest = (file_path/category) |
|
dest.mkdir(exist_ok=True, parents=True) |
|
download_images(dest, urls=search_images(f'{search_term} photo', max_images=max_images)) |
|
resize_images(file_path/category, max_size=400, dest=file_path/category) |
|
|
|
|
|
path = Path('seefood') |
|
search_and_populate("hotdog", "hotdog", path, max_images=90) |
|
for o in ['burger', 'sandwich', 'fruit', 'chips', 'salad']: |
|
search_and_populate(o, "not_hotdog", path, max_images=30) |
|
|
|
failed = verify_images(get_image_files(path)) |
|
failed.map(Path.unlink) |
|
print(f"{len(failed)} failed images") |
|
|
|
dls = DataBlock( |
|
blocks=(ImageBlock, CategoryBlock), |
|
get_items=get_image_files, |
|
splitter=RandomSplitter(valid_pct=0.2), |
|
get_y=parent_label, |
|
item_tfms=[Resize(256, ResizeMethod.Squish)] |
|
).dataloaders(path, bs=32) |
|
|
|
learn = vision_learner(dls, resnet18, metrics=error_rate) |
|
learn.fine_tune(3) |
|
|
|
learn.export("hotdogModel.pkl") |
|
|