KichangKim / DeepDanbooru

AI based multi-label girl image classification system, implemented by using TensorFlow.
MIT License
2.65k stars 260 forks source link

How to use this in a python script on a batch of local images? #56

Closed Bewinxed closed 2 years ago

Bewinxed commented 2 years ago

Hello! Thanks for writing this amazing app, I'm trying to use this to tag a lot of images in a folder on my PC.

Is there any built-in method to embed tags into the images?

I tried using this in a python script but when I try evaluating an image it doesn't return a list of tags or something that I can use.

I'm hoping to use this in a function, that returns a list of tags or a dict, so that I can later embed them into the images.

Any pointers would be greatly appreciated!

rachmadaniHaryono commented 2 years ago

https://gitgud.io/koto/hydrus-dd/-/blob/master/hydrus_dd/evaluate.py#L19-36

that is only to get list of tags, you should implement your own method to embed tag to image

Bewinxed commented 2 years ago

https://gitgud.io/koto/hydrus-dd/-/blob/master/hydrus_dd/evaluate.py#L19-36

that is only to get list of tags, you should implement your own method to embed tag to image

Thanks for the swift response, I'm getting "str" object has no attribute 'input_shape'

Exception has occurred: AttributeError       (note: full exception trace is shown but execution is paused at: <module>)
'str' object has no attribute 'input_shape'
  File "D:\Photos\Art Training (Personal)\classifier.py", line 34, in eval
    for tag, score in tag_sets:
  File "D:\Photos\Art Training (Personal)\classifier.py", line 43, in <module> (Current frame)

This is the code:


projectFolder = 'D:/Model/'
modelFile = r"D:\Model\model-resnet_custom_v3.h5"
tagsFile = r"D:\Model\tags.txt"

taggies = eval("./1011025_10151725838384260_453066707_n.jpg", 0, True, modelFile, tagsFile)
print(taggies)```
rachmadaniHaryono commented 2 years ago

you have to get model and tags like this https://gitgud.io/koto/hydrus-dd/-/blob/master/hydrus_dd/__main__.py#L53-61

Bewinxed commented 2 years ago

you have to get model and tags like this https://gitgud.io/koto/hydrus-dd/-/blob/master/hydrus_dd/__main__.py#L53-61

Thank you so much :) I've ended up creating the following which tags a whole folder, or folder of folders.

Now my library is sfw and tidy, Thanks <3

import pathlib
import asyncio
import os
from libxmp import XMPFiles, XMPMeta, consts
from PIL import Image
import piexif
import io
from typing import Any, List, Optional, Sequence, Tuple, Union
import click
import six
import deepdanbooru

try:
    import tensorflow as tf
except ImportError:
    print("Tensorflow Import failed")
    tf = None

def load_tags(tags_path: Union[pathlib.Path, str, click.types.Path]):
    with open(tags_path, "r") as stream:  # type: ignore
        tags = [tag for tag in (tag.strip() for tag in stream) if tag]
    return tags

def eval(
    image_path: Union[six.BytesIO, str, click.types.Path],
    threshold: float,
    return_score: bool = False,
    model: Optional[Any] = None,
    tags: Optional[List[str]] = None,
) -> Sequence[Union[str, Tuple[str, Any], None]]:

    result_tags = []

    tag_sets = deepdanbooru.commands.evaluate_image(image_path, model, tags, threshold)
    for tag, score in tag_sets:
        if score >= threshold:
            if not return_score:
                result_tags.append(tag)
            else:
                result_tags.append((tag, score))  # type: ignore

    return result_tags

projectFolder = "/home/bewinxed/Photos/gPhotos/Model/"
modelFile = "/home/bewinxed/Photos/gPhotos/Model/model-resnet_custom_v3.h5"
tagsFile = "/home/bewinxed/Photos/gPhotos/Model/tags.txt"

def load_model_and_tags(model_path, tags_path, compile_):
    print("loading model...")
    if compile_ is None:
        model = tf.keras.models.load_model(model_path)
    else:
        model = tf.keras.models.load_model(model_path, compile=compile_)
    print("loading tags...")
    tags = load_tags(tags_path)
    return model, tags

def replace_exif(file, tags):
    """
    Replaces the createdTime Exif data of image with date parsed from the datetime_object
    This function replaces the following three properties of the Exif data of
    the image file that is passed to it:
            Exif.Image.DateTime, at offset 306
            Exif.Photo.DateTimeOriginal at offset 36867
            Exif.Photo.DateTimeDigitized at offset 36868
    This function expects that both PIL and piexif are imported
    For more information about Exif and offsets, see http://www.exiv2.org/tags.html
    Parameters
    ----------
    file : string
            Absolute or relative path to the relevant image file of which the Exif data should
            be changed
    datetime_object : datetime object
            A datetime object generated by the standard datetime module which
            contains the date to which the image's Exif data should be changed

    """

    img = Image.open(file)

    try:
        # If the image already contains data, we only replace the relevant properties
        exif_dict = piexif.load(img.info["exif"])
        print(f"Exif load for file '{file}'' successful")
    except KeyError:
        # If the image has no Exif data, we create the relevant properties
        print(f"No Exif data for file '{file}', creating Exif data instead...")
        exif_dict = {}
        exif_dict["0th"] = {}
        exif_dict["Exif"] = {}

    # We now have a useful Exif dict, time to adjust the values
    d = str(tags).replace("[", "").replace("]", "").replace("'", "")
    exif_dict["Exif"][37510] = d.encode("utf-8")
    # exif_dict["Exif"][40094] = bytes(d, "utf-8")

    # Convert into bytes and dump into file
    exif_bytes = piexif.dump(exif_dict)
    piexif.insert(exif_bytes, file)

    print(f"Exif data replacement for file '{file}' successful")

def add_xmp(file, tags):
    # Write XMP Metadata to image
    xmpfile = XMPFiles(file_path=file, open_forupdate=True)
    xmp = xmpfile.get_xmp()
    # if image has no xmp data, create one
    if xmp is None:
        xmp = XMPMeta()
    # write the tags
    for each in tags:
        # check whether XMP includes 'subject' property,
        # if not, create a new one
        if not xmp.does_property_exist(consts.XMP_NS_DC, "subject"):
            xmp.append_array_item(
                consts.XMP_NS_DC,
                "subject",
                each,
                {"prop_array_is_ordered": True, "prop_value_is_array": True},
            )
        # check whether tag has been written to file
        if not xmp.does_array_item_exist(consts.XMP_NS_DC, "subject", each):
            xmp.append_array_item(consts.XMP_NS_DC, "subject", each)
    if xmpfile.can_put_xmp(xmp):
        xmpfile.put_xmp(xmp)
        xmpfile.close_file()
        return 0
    else:
        print("Unable to write XMP data!")
        return -1

# img = "/home/bewinxed/Photos/gPhotos/Photos/Art Training (Personal)/1_27_14 - 1.jpg"

async def taggie(img, model, tags):

    try:
        taggies = eval(
            img,
            0.7,
            True,
            model,
            tags,
        )
    except:
        return

    taggies = [x[0] for x in taggies]

    try:
        replace_exif(img, taggies)
        add_xmp(img, taggies)
    except:
        None

async def tagFolder(folder, model, tags):
    # model, tags = load_model_and_tags(modelFile, tagsFile, None)

    futures = [
        taggie(os.path.join(folder, filename), model, tags)
        for filename in os.listdir(folder)
        if filename.lower().endswith((".jpg", ".png", ".webp", ".jpeg"))
    ]
    await asyncio.gather(*futures)

async def main():
    masterFolder = "/home/bewinxed/Photos/gPhotos/Photos"
    model, tags = load_model_and_tags(modelFile, tagsFile, None)
    futures = [
        tagFolder(os.path.join(masterFolder, filename), model, tags)
        for filename in os.listdir(masterFolder)
    ]
    await asyncio.gather(*futures)

asyncio.run(main())
# print(taggies)