selenium_image_scraper.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from selenium import webdriver\n",
    "from fastai.vision import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fetch_image_urls(query:str, max_links_to_fetch:int, wd:webdriver, sleep_between_interactions:int=1.5):\n",
    "    def scroll_to_end(wd):\n",
    "        wd.execute_script(\"window.scrollTo(0, document.body.scrollHeight);\")\n",
    "        time.sleep(sleep_between_interactions)    \n",
    "    \n",
    "    # build the google query\n",
    "    search_url = \"https://www.google.com/search?safe=off&site=&tbm=isch&source=hp&q={q}&oq={q}&gs_l=img\"\n",
    "\n",
    "    # load the page\n",
    "    wd.get(search_url.format(q=query))\n",
    "\n",
    "    image_urls = set()\n",
    "    image_count = 0\n",
    "    results_start = 0\n",
    "    while image_count < max_links_to_fetch:\n",
    "        scroll_to_end(wd)\n",
    "\n",
    "        # get all image thumbnail results\n",
    "        thumbnail_results = wd.find_elements_by_css_selector(\"img.Q4LuWd\")\n",
    "        number_results = len(thumbnail_results)\n",
    "        \n",
    "        print(f\"Found: {number_results} search results. Extracting links\"+\n",
    "              f\"from {results_start}:{max_links_to_fetch}\")\n",
    "        \n",
    "        for img in thumbnail_results[results_start:number_results]:\n",
    "            # try to click every thumbnail such that we can get the real image behind it\n",
    "            try:\n",
    "                img.click()\n",
    "                time.sleep(sleep_between_interactions)\n",
    "            except Exception:\n",
    "                continue\n",
    "\n",
    "            # extract image urls    \n",
    "            actual_images = wd.find_elements_by_css_selector('img.n3VNCb')\n",
    "            for actual_image in actual_images:\n",
    "                if actual_image.get_attribute('src') and 'http' in actual_image.get_attribute('src'):\n",
    "                    image_urls.add(actual_image.get_attribute('src'))\n",
    "\n",
    "            image_count = len(image_urls)\n",
    "\n",
    "            if len(image_urls) >= max_links_to_fetch:\n",
    "                print(f\"Found: {len(image_urls)} image links, done!\")\n",
    "                break\n",
    "        else:\n",
    "            print(\"Found:\", len(image_urls), \"image links, looking for more ...\")\n",
    "            time.sleep(30)\n",
    "            return\n",
    "            load_more_button = wd.find_element_by_css_selector(\".mye4qd\")\n",
    "            if load_more_button:\n",
    "                wd.execute_script(\"document.querySelector('.mye4qd').click();\")\n",
    "\n",
    "        # move the result startpoint further down\n",
    "        results_start = len(thumbnail_results)\n",
    "\n",
    "    return image_urls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "browser = webdriver.Firefox()\n",
    "searches = ['search1',\n",
    "            'search2',\n",
    "            'searchN'] # put your searches here\n",
    "path = Path('data') # change this to suit your preferences\n",
    "\n",
    "path.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "for search in searches:\n",
    "    tmp = fetch_image_urls(search,10, browser)\n",
    "    df = pd.DataFrame(tmp)\n",
    "    df.to_csv(path/f'{search}.csv', index=None)\n",
    "    \n",
    "browser.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# optional if you don't want the class names to match the queries\n",
    "# classes = ['class1','class2','classN'] \n",
    "\n",
    "for c in searches:\n",
    "    folder = c\n",
    "    file = c + '.csv'\n",
    "    \n",
    "    dest = path/folder\n",
    "    dest.mkdir(parents=True, exist_ok=True)\n",
    "    \n",
    "    download_images(path/file, dest, max_pics=200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# removes any non-image files\n",
    "for c in searches:\n",
    "    print(c)\n",
    "    verify_images(path/c, delete=True, max_size=500)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}