92 lines
3.1 KiB
Python
92 lines
3.1 KiB
Python
import hashlib
|
|
import io
|
|
import os
|
|
import time
|
|
import requests
|
|
|
|
from PIL import Image
|
|
from selenium import webdriver
|
|
from selenium.webdriver.common.by import By
|
|
|
|
|
|
def fetch_image_urls(query: str, max_links_to_fetch: int, wd: webdriver, sleep_between_interactions: int = 1):
|
|
def scroll_to_end(wd):
|
|
wd.execute_script("window.scrollTo(0, document.body.scrollHeight);")
|
|
time.sleep(sleep_between_interactions)
|
|
|
|
search_url = "https://www.google.com/search?safe=off&site=&tbm=isch&source=hp&q={q}&oq={q}&gs_l=img"
|
|
|
|
wd.get(search_url.format(q=query))
|
|
|
|
image_urls = set()
|
|
image_count = 0
|
|
results_start = 0
|
|
|
|
while image_count < max_links_to_fetch:
|
|
scroll_to_end(wd)
|
|
thumbnail_results = wd.find_elements(by=By.CSS_SELECTOR, value="img.Q4LuWd")
|
|
number_results = len(thumbnail_results)
|
|
|
|
print(f"Found: {number_results} search results. Extracting links from {results_start}:{number_results}")
|
|
|
|
for img in thumbnail_results[results_start:number_results]:
|
|
try:
|
|
img.click()
|
|
time.sleep(sleep_between_interactions)
|
|
except Exception:
|
|
continue
|
|
|
|
actual_images = wd.find_elements(by=By.CSS_SELECTOR, value="img.n3VNCb")
|
|
for actual_image in actual_images:
|
|
if actual_image.get_attribute('src') and 'http' in actual_image.get_attribute('src'):
|
|
image_urls.add(actual_image.get_attribute('src'))
|
|
|
|
image_count = len(image_urls)
|
|
|
|
if len(image_urls) >= max_links_to_fetch:
|
|
print(f"Found: {len(image_urls)} image links, done!")
|
|
break
|
|
else:
|
|
print("Found:", len(image_urls), "image links, looking for more ...")
|
|
time.sleep(1)
|
|
results_start = len(thumbnail_results)
|
|
|
|
return image_urls
|
|
|
|
def download_image(folder_path:str, url:str):
|
|
try:
|
|
image_content = requests.get(url).content
|
|
|
|
except Exception as e:
|
|
print(f"ERROR - Could not download {url} - {e}")
|
|
|
|
try:
|
|
image_file = io.BytesIO(image_content)
|
|
image = Image.open(image_file).convert('RGB')
|
|
file_path = os.path.join(folder_path,hashlib.sha1(image_content).hexdigest()[:10] + '.jpg')
|
|
with open(file_path, 'wb') as f:
|
|
image.save(f, "JPEG", quality=100)
|
|
print(f"SUCCESS - saved {url} - as {file_path}")
|
|
except Exception as e:
|
|
print(f"ERROR - Could not save {url} - {e}")
|
|
|
|
def search_and_download(search_term:str, target_path='../learning/train', number_images=int):
|
|
target_folder = os.path.join(target_path, '_'.join(search_term.lower().split(' ')))
|
|
|
|
if not os.path.exists(target_folder):
|
|
os.makedirs(target_folder)
|
|
|
|
with webdriver.Chrome() as wd:
|
|
res = fetch_image_urls(search_term, number_images, wd=wd, sleep_between_interactions=0.5)
|
|
|
|
for elem in res:
|
|
download_image(target_folder, elem)
|
|
|
|
print("Input search term: ", end='')
|
|
search_term = input()
|
|
print("Input number of images to download: ", end='')
|
|
number_images = int(input())
|
|
|
|
search_and_download(search_term=search_term, number_images=number_images)
|
|
|