Tagging Anime Images in Lychee with DeepDanbooru

I wanted to create a collection of my favorite anime images (e.g. Pixiv artworks) on my private Lychee instance. Uploading the images was an easy task, but I was thinking that the browsing experience would be better if there was a method to filter the images by their characteristics. Fortunately, I was able to find a project on GitHub named DeepDanbooru.

GitHub - KichangKim/DeepDanbooru: AI based multi-label girl image classification system, implemented by using TensorFlow.

AI based multi-label girl image classification system, implemented by using TensorFlow. - GitHub - KichangKim/DeepDanbooru: AI based multi-label girl image classification system, implemented by usi…

GitHubKichangKim

Lychee supports adding tags to images. The plan was to tag the images using this project, and the search box in Lychee already supports searching images by tags. Using the Lychee API, the images will be fetched and passed into the AI model. Tags that are outputted will then be added to the corresponding images.

To use DeepDanbooru, we first clone the git repository. Install the Python dependencies using pip install -r requirements.txt. Then, we can download a pre-trained model from the Releases page. You may also train your own model by following the instructions in the repo's README. We then place the downloaded folder in the repo's root directory and rename it to model.

Next, create a script folder in which we will place our automation script. Open the folder. Copy the content below to a file named requirements.txt.

alive_progress
halo
titlecase
yachalk

Run pip install -r requirements.txt. Copy the content below to a file named script.py.

import os
import requests
from alive_progress import alive_bar
from deepdanbooru.commands import evaluate_image
from deepdanbooru.project import load_model_from_project, load_tags_from_project
from halo import Halo
from io import BytesIO
from titlecase import titlecase
from yachalk import chalk


#
# Configuration
#

LYCHEE_HOST: str = "http://localhost"
"""
Host of the Lychee instance.
We recommend running this script on the same server hosting the Lychee instance for the best performance and the least network transfer.
Defaults to `http://localhost`.
"""
LYCHEE_TOKEN: str = ""
"""
API token for the Lychee instance.
"""

ALBUMS_LIST_RULE: str = "exclude"  # "include" | "exclude"
"""
Rule for the albums list.
Defaults to `exclude`.
"""

ALBUMS_LIST: list[str] = []
"""
IDs of albums to include or exclude (depending on `ALBUMS_LIST_RULE`).
"""

MODEL_PATH: str = os.path.abspath("../model")
"""
Path to the DeepDanbooru model directory.
Defaults to `os.path.abspath("../model")`, which assumes that you have placed this script in a subfolder under the DeepDanbooru folder.
"""

TAG_THRESHOLD: float = 0.5
"""
Threshold for applying tags. Tags under this threshold will be filtered out.
Defaults to `0.5`.
"""

TAG_TAGGED: str = "Tagged by DeepDanbooru"
"""
Raw tag (not transformed by `CONVERT_TAG`) for marking photos that have been already processed.
Photos that already have this tag will be skipped.
Defaults to `Tagged by DeepDanbooru`.
"""

ADDITIONAL_TAGS: list[str] = ["Source: Pixiv"]
"""
Additional raw tags (not transformed by `CONVERT_TAG`) to be added to the photos.
"""

OVERRIDE_EXISTING_TAGS: bool = False
"""
Whether to override or append existing tags using new generated tags.
Defaults to `True`.
"""


def CONVERT_TAG(tag: str) -> str:
    """
    Custom function to transform a tag before storing it in Lychee.
    """
    if tag.startswith("rating:"):
        return "Rating: " + titlecase(tag[7:])
    converted_tag = titlecase(tag.replace("_", " "))
    return converted_tag


#
# Constants
#

REPLACE: str = "\033[1A\033[2K"


#
# Main
#

# Load model and tags
spinner: Halo = Halo(text=chalk.gray("Loading model..."), color="grey").start()
dd_model = load_model_from_project(MODEL_PATH, compile_model=False)
spinner.succeed(chalk.green("Model loaded successfully."))
spinner: Halo = Halo(text=chalk.gray("Loading tags..."), color="grey").start()
dd_tags = load_tags_from_project(MODEL_PATH)
spinner.succeed(chalk.green("Tags loaded successfully."))
print()

# Define headers
headers: dict = {
    "Accept": "application/json",
    "Authorization": LYCHEE_TOKEN,
    "Content-Type": "application/json",
}

# Fetch albums from Lychee
spinner: Halo = Halo(text=chalk.gray("Fetching albums..."), color="grey").start()
response: requests.Response = requests.post(
    f"{LYCHEE_HOST}/api/Albums::get", json={}, headers=headers
)
all_albums: list[dict] = response.json()["albums"]
albums: list[dict] = [
    album
    for album in all_albums
    if (
        album["id"] in ALBUMS_LIST
        if ALBUMS_LIST_RULE == "include"
        else album["id"] not in ALBUMS_LIST
    )
]
spinner.succeed(
    f"Found {len(all_albums)} albums in total, {len(albums)} albums to process."
)

# Process each album
for album in albums:

    def print_separator():
        term_size = os.get_terminal_size()
        print("=" * term_size.columns)

    print()
    print_separator()
    print(
        chalk.cyan.bold(f"ALBUM {chalk.underline(album['title'])} (ID: {album['id']})")
    )

    # Fetch photos from the album
    spinner: Halo = Halo(text=chalk.gray("Fetching photos..."), color="grey").start()
    response: requests.Response = requests.post(
        f"{LYCHEE_HOST}/api/Album::get", json={"albumID": album["id"]}, headers=headers
    )
    photos: list[dict] = response.json()["photos"]
    spinner.succeed(f"Found {len(photos)} photos.")

    # Process each photo
    with alive_bar(len(photos)) as bar:
        for photo in photos:
            # Skip if the photo has already been processed
            if TAG_TAGGED in photo["tags"]:
                print(
                    chalk.gray(
                        f"PHOTO {chalk.underline(photo['title'])} (ID: {photo['id']}) (skipped)"
                    )
                )
                bar()
                continue

            print(
                chalk.cyan(
                    f"PHOTO {chalk.underline(photo['title'])} (ID: {photo['id']})"
                )
            )

            # Download the image
            spinner: Halo = Halo(
                text=chalk.gray("Downloading image..."), color="grey"
            ).start()
            image_data: bytes = requests.get(
                f"{LYCHEE_HOST}/{photo['size_variants']['original']['url']}"
            ).content
            spinner.stop()

            # Process the image with DeepDanbooru
            result: list[tuple[str, float]] = evaluate_image(
                BytesIO(image_data), dd_model, dd_tags, TAG_THRESHOLD
            )

            # Convert tags
            tags: list[str] = []
            tags.extend(ADDITIONAL_TAGS)
            tags.append(TAG_TAGGED)
            tags.extend([CONVERT_TAG(tag) for tag, score in result])
            print(chalk.blue(f"Tags: {', '.join(tags)}"))

            # Upload tags
            spinner: Halo = Halo(
                text=chalk.gray("Uploading tags..."), color="grey"
            ).start()
            response: requests.Response = requests.post(
                f"{LYCHEE_HOST}/api/Photo::setTags",
                json={
                    "photoIDs": [photo["id"]],
                    "tags": tags,
                    "shall_override": OVERRIDE_EXISTING_TAGS,
                },
                headers=headers,
            )
            if response.status_code == 204:
                spinner.succeed(chalk.green("Uploaded successfully."))
            else:
                spinner.fail(chalk.red("Failed to upload."))

            bar()

    print_separator()

The script is self-explanatory. Configure the values according to your instance and your preferences. Run the script and see things happen.