booru-i2i-scraper.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard",
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Linaqruf/sd-notebook-collection/blob/main/booru-i2i-scraper.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "![](https://visitor-badge.glitch.me/badge?page_id=linaqruf.i2i-scraper) [![GitHub](https://badgen.net/badge/icon/github?icon=github&label)](https://github.com/Linaqruf/sd-notebook-collection/blob/main/booru-i2i-scraper.ipynb) \n",
        "\n",
        "# **Booru I2I Scraper**\n",
        "Not perfect but usable<br>\n",
        "\n",
        "<details>\n",
        "  <summary><big>Support Us!</big></summary>\n",
        "    <ul>\n",
        "      <li>\n",
        "        <a href=\"https://ko-fi.com/linaqruf\">\n",
        "          <img src=\"https://img.shields.io/badge/Support%20me%20on%20Ko--fi-F16061?logo=ko-fi&logoColor=white&style=flat\" alt=\"Ko-fi badge\">\n",
        "        </a>\n",
        "      </li>\n",
        "      <li>\n",
        "        <a href=\"https://saweria.co/linaqruf\">\n",
        "          <img src=\"https://img.shields.io/badge/Saweria-7B3F00?style=flat&logo=ko-fi&logoColor=white\" alt=\"Saweria badge\">\n",
        "        </a>\n",
        "      </li>\n",
        "    </ul>\n",
        "</details>"
      ],
      "metadata": {
        "id": "3aKyBtukgWx3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Install Dependencies\n",
        "import os\n",
        "from IPython.utils import capture\n",
        "\n",
        "print(\"\u001b[1;32mInstalling...\")\n",
        "print(\"\u001b[1;32mPlease wait...\\n\")\n",
        "with capture.capture_output() as cap:\n",
        "  !pip -q install opencv-python tensorflow faiss-gpu huggingface_hub \"gradio==3.16.2\" gallery-dl\n",
        "  !pip -q uninstall -y Pillow\n",
        "  !pip -q install \"pillow>=9.1.0\"\n",
        "  del cap\n",
        "print(\"\u001b[1;32mDone! Restarting...\")\n",
        "\n",
        "os.kill(os.getpid(), 9)"
      ],
      "metadata": {
        "id": "7OaAW3Bt9m6S",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Launch\n",
        "import argparse\n",
        "import functools\n",
        "import json\n",
        "import os \n",
        "import zipfile\n",
        "import shutil\n",
        "from IPython.utils import capture\n",
        "from tqdm import tqdm\n",
        "from pathlib import Path\n",
        "from tensorflow.keras.models import load_model\n",
        "import faiss\n",
        "import PIL.Image\n",
        "import gradio as gr\n",
        "import numpy as np\n",
        "\n",
        "import requests\n",
        "import tensorflow as tf\n",
        "from huggingface_hub import hf_hub_download\n",
        "\n",
        "root_dir = \"/content\"\n",
        "dir = os.path.join(root_dir, \"danbooru2022_image_similarity\")\n",
        "repo_url = \"https://huggingface.co/spaces/SmilingWolf/danbooru2022_image_similarity\"\n",
        "utils_dir = os.path.join(dir, \"Utils\")\n",
        "index_dir = os.path.join(dir, \"index\")\n",
        "deps_dir = os.path.join(root_dir, \"deps\")\n",
        "image_dir = os.path.join(root_dir, \"images\")\n",
        "app_py = os.path.join(repo_url, \"resolve/main/app.py\")\n",
        "index_dir_files = [os.path.join(repo_url, \"resolve/main/index/cosine_ids.npy\"),\n",
        "                   os.path.join(repo_url, \"resolve/main/index/cosine_infos.json\"),\n",
        "                   os.path.join(repo_url, \"resolve/main/index/cosine_knn.index\")]\n",
        "utils = os.path.join(repo_url, \"resolve/main/Utils/dbimutils.py\")\n",
        "\n",
        "TITLE = \"## Danbooru Explorer\"\n",
        "DESCRIPTION = \"\"\"\n",
        "Image similarity-based retrieval tool using:\n",
        "- [SmilingWolf/wd-v1-4-convnext-tagger-v2](https://huggingface.co/SmilingWolf/wd-v1-4-convnext-tagger-v2) as feature extractor\n",
        "- [Faiss](https://github.com/facebookresearch/faiss) and [autofaiss](https://github.com/criteo/autofaiss) for indexing\n",
        "\"\"\"\n",
        "CONV_FEXT_LAYER = \"predictions_norm\"\n",
        "\n",
        "files = [\"keras_metadata.pb\", \"saved_model.pb\", \"selected_tags.csv\"]\n",
        "sub_dir = \"variables\"\n",
        "sub_dir_files = [\"variables.data-00000-of-00001\", \"variables.index\"]\n",
        "csv_file = files[-1]\n",
        "model_repo = \"SmilingWolf/wd-v1-4-convnext-tagger-v2\"\n",
        "model_dir = \"/content/wd14_tagger\"\n",
        "\n",
        "for directory in [dir, utils_dir, index_dir, deps_dir, image_dir]:\n",
        "  os.makedirs(directory, exist_ok=True)\n",
        "\n",
        "def ubuntu_deps(url, name, dst):\n",
        "  os.makedirs(dst, exist_ok=True)\n",
        "  !wget -q --show-progress {url}\n",
        "  with zipfile.ZipFile(name, 'r') as deps:\n",
        "    deps.extractall(dst)\n",
        "  !dpkg -i {dst}/*\n",
        "  os.remove(name)\n",
        "  shutil.rmtree(dst)\n",
        "\n",
        "print(\"\u001b[1;32mInstalling...\")\n",
        "with capture.capture_output() as cap:\n",
        "  !apt -y update -qq\n",
        "  ubuntu_deps(\"https://huggingface.co/Linaqruf/fast-repo/resolve/main/ram_patch.zip\", \"ram_patch.zip\", deps_dir)\n",
        "  %env LD_PRELOAD=libtcmalloc.so\n",
        "  ubuntu_deps(\"https://huggingface.co/Linaqruf/fast-repo/resolve/main/deb-libs.zip\", \"deb-libs.zip\", deps_dir)\n",
        "  del cap\n",
        "\n",
        "print(\"\u001b[1;32mDownloading Danbooru Explorer...\")\n",
        "with capture.capture_output() as cap:\n",
        "  for file in index_dir_files:\n",
        "    index_basename = os.path.basename(file)\n",
        "    !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {index_dir} -o {index_basename} {file}\n",
        "  !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {dir} -o \"app.py\" {app_py}\n",
        "  !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {utils_dir} -o \"dbimutils.py\" {utils}\n",
        "  del cap\n",
        "\n",
        "init = \"\"\"\"\"\"\n",
        "with open(os.path.join(utils_dir, \"__init__.py\"), 'w') as f:\n",
        "    f.write(init)\n",
        "\n",
        "setup = \"\"\"from setuptools import setup, find_packages\n",
        "setup(name = \"Utils\", packages = find_packages())\"\"\"\n",
        "with open(os.path.join(dir, \"setup.py\"), 'w') as f:\n",
        "    f.write(setup)\n",
        "\n",
        "os.chdir(dir)\n",
        "print(\"\u001b[1;32mInstall setup.py...\")\n",
        "with capture.capture_output() as cap:\n",
        "  !pip -q install .\n",
        "  del cap\n",
        "\n",
        "import Utils.dbimutils as dbimutils\n",
        "\n",
        "if not os.path.exists(model_dir):\n",
        "  print(f\"\u001b[1;32mDownloading WD 14 Tagger model from {model_repo}...\")\n",
        "  with capture.capture_output() as cap:\n",
        "    for file in files:\n",
        "      hf_hub_download(model_repo, file, cache_dir=model_dir, force_download=True, force_filename=file)\n",
        "    for file in sub_dir_files:\n",
        "      hf_hub_download(model_repo, file, subfolder=sub_dir, cache_dir=os.path.join(\n",
        "          model_dir, sub_dir), force_download=True, force_filename=file)\n",
        "    del cap\n",
        "    \n",
        "def load_model(model_path, feature_extraction_layer):\n",
        "    full_model = tf.keras.models.load_model(model_path)\n",
        "    model = tf.keras.models.Model(\n",
        "        full_model.inputs, full_model.get_layer(feature_extraction_layer).output\n",
        "    )\n",
        "    return model\n",
        "\n",
        "def danbooru_id_to_url(image_id, selected_ratings, api_username=\"\", api_key=\"\"):\n",
        "    headers = {\"User-Agent\": \"image_similarity_tool\"}\n",
        "    ratings_to_letters = {\n",
        "        \"General\": \"g\",\n",
        "        \"Sensitive\": \"s\",\n",
        "        \"Questionable\": \"q\",\n",
        "        \"Explicit\": \"e\",\n",
        "    }\n",
        "\n",
        "    acceptable_ratings = [ratings_to_letters[x] for x in selected_ratings]\n",
        "\n",
        "    image_url = f\"https://danbooru.donmai.us/posts/{image_id}.json\"\n",
        "    if api_username != \"\" and api_key != \"\":\n",
        "        image_url = f\"{image_url}?api_key={api_key}&login={api_username}\"\n",
        "\n",
        "    r = requests.get(image_url, headers=headers)\n",
        "    if r.status_code != 200:\n",
        "        return None\n",
        "\n",
        "    content = json.loads(r.text)\n",
        "    image_url = content[\"large_file_url\"] if \"large_file_url\" in content else None\n",
        "    image_url = image_url if content[\"rating\"] in acceptable_ratings else None\n",
        "    return image_url\n",
        "\n",
        "class SimilaritySearcher:\n",
        "    def __init__(self, model, images_ids):\n",
        "        self.knn_index = None\n",
        "        self.knn_metric = None\n",
        "\n",
        "        self.model = model\n",
        "        self.images_ids = images_ids\n",
        "\n",
        "    def change_index(self, knn_metric):\n",
        "        if knn_metric == self.knn_metric:\n",
        "            return\n",
        "\n",
        "        if knn_metric == \"ip\":\n",
        "            self.knn_index = faiss.read_index(\"index/ip_knn.index\")\n",
        "            config = json.loads(open(\"index/ip_infos.json\").read())[\"index_param\"]\n",
        "        elif knn_metric == \"cosine\":\n",
        "            self.knn_index = faiss.read_index(\"index/cosine_knn.index\")\n",
        "            config = json.loads(open(\"index/cosine_infos.json\").read())[\"index_param\"]\n",
        "\n",
        "        faiss.ParameterSpace().set_index_parameters(self.knn_index, config)\n",
        "        self.knn_metric = knn_metric\n",
        "\n",
        "\n",
        "    def predict(\n",
        "        self, image, selected_ratings, knn_metric, api_username, api_key, n_neighbours\n",
        "    ):\n",
        "        _, height, width, _ = self.model.inputs[0].shape\n",
        "\n",
        "        self.change_index(knn_metric)\n",
        "\n",
        "        # Alpha to white\n",
        "        image = image.convert(\"RGBA\")\n",
        "        new_image = PIL.Image.new(\"RGBA\", image.size, \"WHITE\")\n",
        "        new_image.paste(image, mask=image)\n",
        "        image = new_image.convert(\"RGB\")\n",
        "        image = np.asarray(image)\n",
        "\n",
        "        # PIL RGB to OpenCV BGR\n",
        "        image = image[:, :, ::-1]\n",
        "\n",
        "        image = dbimutils.make_square(image, height)\n",
        "        image = dbimutils.smart_resize(image, height)\n",
        "        image = image.astype(np.float32)\n",
        "        image = np.expand_dims(image, 0)\n",
        "        target = self.model(image).numpy()\n",
        "\n",
        "        if self.knn_metric == \"cosine\":\n",
        "            faiss.normalize_L2(target)\n",
        "\n",
        "        dists, indexes = self.knn_index.search(target, k=n_neighbours)\n",
        "        neighbours_ids = self.images_ids[indexes][0]\n",
        "        neighbours_ids = [int(x) for x in neighbours_ids]\n",
        "\n",
        "        captions = []\n",
        "        for image_id, dist in zip(neighbours_ids, dists[0]):\n",
        "            captions.append(f\"{image_id}/{dist:.2f}\")\n",
        "\n",
        "      \n",
        "        image_urls = []\n",
        "        scraper_text = os.path.join(root_dir, \"scrape_this.txt\")\n",
        "        \n",
        "        for image_id in neighbours_ids:\n",
        "          current_url = danbooru_id_to_url(\n",
        "            image_id, selected_ratings, api_username, api_key\n",
        "          )\n",
        "          if current_url is not None:              \n",
        "            image_urls.append(current_url)\n",
        "\n",
        "        with open(scraper_text, 'w') as f:\n",
        "          f.write('\\n'.join(image_urls))\n",
        "\n",
        "        !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {image_dir} -i {scraper_text}\n",
        "\n",
        "        return list(zip(image_urls, captions))\n",
        "\n",
        "\n",
        "def main():\n",
        "    gpus = tf.config.list_physical_devices('GPU')\n",
        "    if len(gpus) > 0:\n",
        "        tf.config.set_visible_devices(gpus[0], 'GPU')\n",
        "        tf.config.experimental.set_memory_growth(gpus[0], True)\n",
        "    else:\n",
        "        print(\"No GPU found, using CPU instead.\")\n",
        "    model = load_model(model_dir, CONV_FEXT_LAYER)\n",
        "    images_ids = np.load(\"index/cosine_ids.npy\")\n",
        "\n",
        "    searcher = SimilaritySearcher(model=model, images_ids=images_ids)\n",
        "\n",
        "    with gr.Blocks() as demo:\n",
        "        gr.Markdown(TITLE)\n",
        "        gr.Markdown(DESCRIPTION)\n",
        "\n",
        "        with gr.Row():\n",
        "            input = gr.Image(type=\"pil\", label=\"Input\")\n",
        "            with gr.Column():\n",
        "                with gr.Row():\n",
        "                    api_username = gr.Textbox(label=\"Danbooru API Username\")\n",
        "                    api_key = gr.Textbox(label=\"Danbooru API Key\")\n",
        "                selected_ratings = gr.CheckboxGroup(\n",
        "                    choices=[\"General\", \"Sensitive\", \"Questionable\", \"Explicit\"],\n",
        "                    value=[\"General\", \"Sensitive\"],\n",
        "                    label=\"Ratings\",\n",
        "                )\n",
        "                with gr.Row():\n",
        "                    selected_metric = gr.Radio(\n",
        "                        choices=[\"cosine\"],\n",
        "                        value=\"cosine\",\n",
        "                        label=\"Metric selection\",\n",
        "                        visible=False,\n",
        "                    )\n",
        "                    n_neighbours = gr.Slider(\n",
        "                        minimum=1, maximum=500, value=5, step=1, label=\"# of images\"\n",
        "                    )\n",
        "                find_btn = gr.Button(\"Find similar images\")\n",
        "        similar_images = gr.Gallery(label=\"Similar images\")\n",
        "\n",
        "        similar_images.style(grid=5)\n",
        "        find_btn.click(\n",
        "            fn=searcher.predict,\n",
        "            inputs=[\n",
        "                input,\n",
        "                selected_ratings,\n",
        "                selected_metric,\n",
        "                api_username,\n",
        "                api_key,\n",
        "                n_neighbours,\n",
        "            ],\n",
        "            outputs=[similar_images],\n",
        "        )\n",
        "\n",
        "    demo.queue()\n",
        "    demo.launch(share=True, debug=True)\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()"
      ],
      "metadata": {
        "id": "F4nNbab7-9tL",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title  ## 📝 Download Images\n",
        "#@markdown Download file manually from files tab or save to Google Drive\n",
        "%cd /content/\n",
        "\n",
        "!zip -r /content/images.zip images\n",
        "\n",
        "from pydrive.auth import GoogleAuth\n",
        "from google.colab import drive\n",
        "from pydrive.drive import GoogleDrive\n",
        "from google.colab import auth\n",
        "from oauth2client.client import GoogleCredentials\n",
        "def create_folder(folder_name):\n",
        "    # Check if folder exists\n",
        "    file_list = drive.ListFile({'q': \"title='{}' and mimeType='application/vnd.google-apps.folder' and trashed=false\".format(folder_name)}).GetList()\n",
        "    if len(file_list) > 0:\n",
        "        # Folder exists\n",
        "        print('Debug: Folder exists')\n",
        "        folder_id = file_list[0]['id']\n",
        "    else:\n",
        "        print('Debug: Creating folder')\n",
        "        file = drive.CreateFile({'title': folder_name, 'mimeType': 'application/vnd.google-apps.folder'})\n",
        "        file.Upload()\n",
        "        folder_id = file.attr['metadata']['id']\n",
        "    # return folder id\n",
        "    return folder_id\n",
        "# Upload file to Google Drive\n",
        "def upload_file(file_name, folder_id, save_as):\n",
        "    # Check if file exists\n",
        "    file_list = drive.ListFile({'q': \"title='{}' and trashed=false\".format(save_as)}).GetList()\n",
        "    if len(file_list) > 0:\n",
        "        print('Debug: File already exists')\n",
        "        # Change file name to avoid overwriting\n",
        "        save_as = save_as + ' (1)'\n",
        "    file = drive.CreateFile({'title': save_as, 'parents': [{'id': folder_id}]})\n",
        "    file.SetContentFile(file_name)\n",
        "    # Upload and set permission to public\n",
        "    file.Upload()\n",
        "    file.InsertPermission({'type': 'anyone', 'value': 'anyone', 'role': 'reader'})\n",
        "    # return file id\n",
        "    return file.attr['metadata']['id']\n",
        "\n",
        "use_drive = True #@param {type:\"boolean\"}\n",
        "folder_name = \"scraper\" #@param {type: \"string\"}\n",
        "save_as = \"nakiriayame.zip\" #@param {type: \"string\"}\n",
        "\n",
        "if use_drive:\n",
        "  auth.authenticate_user()\n",
        "  gauth = GoogleAuth()\n",
        "  gauth.credentials = GoogleCredentials.get_application_default()\n",
        "  drive = GoogleDrive(gauth)\n",
        "  file_id = upload_file('/content/images.zip', create_folder(folder_name), save_as)\n",
        "  print(\"Your sharing link: https://drive.google.com/file/d/\" + file_id + \"/view?usp=sharing\")  "
      ],
      "metadata": {
        "cellView": "form",
        "id": "1gH_zBqtTyEb"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}