-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
68 lines (57 loc) · 2.21 KB
/
main.py
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
from duckduckgo_search import DDGS
from fastcore.all import *
# create venv
# python3 -m venv ./venv
def search_images(term, max_images=80):
print(f"Searching for '{term}'")
with DDGS() as ddgs:
ddgs_images_gen = ddgs.images(term)
count = 0
ddgs_images_list = []
while count < max_images:
image = next(ddgs_images_gen)
ddgs_images_list.append(image.get('image'))
count = count+1
return ddgs_images_list
if __name__ == "__main__":
#NB: `search_images` depends on duckduckgo.com, which doesn't always return correct responses.
# If you get a JSON error, just try running it again (it may take a couple of tries).
urls = search_images('sandwich photos', max_images=1)
urls[0]
from fastdownload import download_url
dest = 'sandwich.jpg'
download_url(urls[0], dest, show_progress=False)
from fastai.vision.all import *
im = Image.open(dest)
im.to_thumb(256,256)
download_url(search_images('peninsula photo', max_images=1)[0], 'peninsula.jpg', show_progress=False)
Image.open('peninsula.jpg').to_thumb(256,256)
searches = 'peninsula-photo','sandwich-photo'
path = Path('peninsula-or-sandwich')
from time import sleep
for o in searches:
dest = (path/o)
dest.mkdir(exist_ok=True, parents=True)
download_images(dest, urls=search_images(f'{o}'))
sleep(10)
resize_images(path/o, max_size=400, dest=path/o)
failed = verify_images(get_image_files(path))
failed.map(Path.unlink)
print(len(failed))
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)
dls.show_batch(max_n=6)
learn = vision_learner(dls, resnet18, metrics=error_rate)
learn.fine_tune(6)
p = 'peninsula.jpg'
is_what,_,probs = learn.predict(PILImage.create(p))
print(f"{p} is a: {is_what}.")
p = 'sandwich.jpg'
is_what,_,probs = learn.predict(PILImage.create(p))
print(f"{p} is a: {is_what}.")
learn.export("peninsula-or-sandwich.pkl")