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