Also-k/DeepDanbooru_string

1 2#!/usr/bin/env python3 4from __future__ import annotations5 6import argparse7import functools8import os9import html10import pathlib11import tarfile12 13import deepdanbooru as dd14import gradio as gr15import huggingface_hub16import numpy as np17import PIL.Image18import tensorflow as tf19import piexif20import piexif.helper21 22TITLE = 'DeepDanbooru String'23 24TOKEN = os.environ['TOKEN']25MODEL_REPO = 'CikeyQI/DeepDanbooru_string'26MODEL_FILENAME = 'model-resnet_custom_v3.h5'27LABEL_FILENAME = 'tags.txt'28 29 30def parse_args() -> argparse.Namespace:31    parser = argparse.ArgumentParser()32    parser.add_argument('--score-slider-step', type=float, default=0.05)33    parser.add_argument('--score-threshold', type=float, default=0.5)34    parser.add_argument('--theme', type=str, default='dark-grass')35    parser.add_argument('--live', action='store_true')36    parser.add_argument('--share', action='store_true')37    parser.add_argument('--port', type=int)38    parser.add_argument('--disable-queue',39                        dest='enable_queue',40                        action='store_false')41    parser.add_argument('--allow-flagging', type=str, default='never')42    return parser.parse_args()43 44 45def load_sample_image_paths() -> list[pathlib.Path]:46    image_dir = pathlib.Path('images')47    if not image_dir.exists():48        dataset_repo = 'hysts/sample-images-TADNE'49        path = huggingface_hub.hf_hub_download(dataset_repo,50                                               'images.tar.gz',51                                               repo_type='dataset',52                                               use_auth_token=TOKEN)53        with tarfile.open(path) as f:54            f.extractall()55    return sorted(image_dir.glob('*'))56 57 58def load_model() -> tf.keras.Model:59    path = huggingface_hub.hf_hub_download(MODEL_REPO,60                                           MODEL_FILENAME,61                                           use_auth_token=TOKEN)62    model = tf.keras.models.load_model(path)63    return model64 65 66def load_labels() -> list[str]:67    path = huggingface_hub.hf_hub_download(MODEL_REPO,68                                           LABEL_FILENAME,69                                           use_auth_token=TOKEN)70    with open(path) as f:71        labels = [line.strip() for line in f.readlines()]72    return labels73 74def plaintext_to_html(text):75    text = "<p>" + "<br>\n".join([f"{html.escape(x)}" for x in text.split('\n')]) + "</p>"76    return text77 78def predict(image: PIL.Image.Image, score_threshold: float,79            model: tf.keras.Model, labels: list[str]) -> dict[str, float]:80    rawimage = image81    _, height, width, _ = model.input_shape82    image = np.asarray(image)83    image = tf.image.resize(image,84                            size=(height, width),85                            method=tf.image.ResizeMethod.AREA,86                            preserve_aspect_ratio=True)87    image = image.numpy()88    image = dd.image.transform_and_pad_image(image, width, height)89    image = image / 255.90    probs = model.predict(image[None, ...])[0]91    probs = probs.astype(float)92    res = dict()93    for prob, label in zip(probs.tolist(), labels):94        if prob < score_threshold:95            continue96        res[label] = prob97    b = dict(sorted(res.items(),key=lambda item:item[1], reverse=True))98    a = ', '.join(list(b.keys())).replace('_',' ').replace('(','\(').replace(')','\)')99    c = ', '.join(list(b.keys()))100    101    items = rawimage.info102    geninfo = ''103    104    if "exif" in rawimage.info:105        exif = piexif.load(rawimage.info["exif"])106        exif_comment = (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b'')107        try:108            exif_comment = piexif.helper.UserComment.load(exif_comment)109        except ValueError:110            exif_comment = exif_comment.decode('utf8', errors="ignore")111    112        items['exif comment'] = exif_comment113        geninfo = exif_comment114    115        for field in ['jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'dpi', 'exif',116                      'loop', 'background', 'timestamp', 'duration']:117            items.pop(field, None)118    119    geninfo = items.get('parameters', geninfo)120    121    info = f"""122<p><h4>PNG Info</h4></p>    123"""124    for key, text in items.items():125        info += f"""126<div>127<p><b>{plaintext_to_html(str(key))}</b></p>128<p>{plaintext_to_html(str(text))}</p>129</div>130""".strip()+"\n"131    132    if len(info) == 0:133        message = "Nothing found in the image."134        info = f"<div><p>{message}<p></div>"135    136    return (a,c,res,info)137 138 139def main():140    args = parse_args()141    model = load_model()142    labels = load_labels()143 144    func = functools.partial(predict, model=model, labels=labels)145    func = functools.update_wrapper(func, predict)146 147    gr.Interface(148        func,149        [150            gr.inputs.Image(type='pil', label='Input'),151            gr.inputs.Slider(0,152                             1,153                             step=args.score_slider_step,154                             default=args.score_threshold,155                             label='Score Threshold'),156        ],157        [158            gr.outputs.Textbox(label='Output (string)'), 159            gr.outputs.Textbox(label='Output (raw string)'), 160            gr.outputs.Label(label='Output (label)'),161            gr.outputs.HTML()162        ],163        examples=[164        ['miku.jpg',0.5],165        ['miku2.jpg',0.5]166        ],167        title=TITLE,168        description='''169Demo for [KichangKim/DeepDanbooru](https://github.com/KichangKim/DeepDanbooru) with "ready to copy" prompt and a prompt analyzer.170 171Modified from [hysts/DeepDanbooru](https://huggingface.co/spaces/hysts/DeepDanbooru)172 173PNG Info code forked from [AUTOMATIC1111/stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui)174        ''',175        theme=args.theme,176        allow_flagging=args.allow_flagging,177        live=args.live,178    ).launch(179        enable_queue=args.enable_queue,180        server_port=args.port,181        share=args.share,182    )183 184 185if __name__ == '__main__':186    main()