Duplicate from hysts/DeepDanbooru

Co-authored-by: hysts <[email protected]>

.gitattributes

@@ -0,0 +1,27 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+images

.pre-commit-config.yaml

@@ -0,0 +1,36 @@
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.2.0
+  hooks:
+  - id: check-executables-have-shebangs
+  - id: check-json
+  - id: check-merge-conflict
+  - id: check-shebang-scripts-are-executable
+  - id: check-toml
+  - id: check-yaml
+  - id: double-quote-string-fixer
+  - id: end-of-file-fixer
+  - id: mixed-line-ending
+    args: ['--fix=lf']
+  - id: requirements-txt-fixer
+  - id: trailing-whitespace
+- repo: https://github.com/myint/docformatter
+  rev: v1.4
+  hooks:
+  - id: docformatter
+    args: ['--in-place']
+- repo: https://github.com/pycqa/isort
+  rev: 5.12.0
+  hooks:
+    - id: isort
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v0.991
+  hooks:
+    - id: mypy
+      args: ['--ignore-missing-imports']
+      additional_dependencies: ['types-python-slugify']
+- repo: https://github.com/google/yapf
+  rev: v0.32.0
+  hooks:
+  - id: yapf
+    args: ['--parallel', '--in-place']

.style.yapf

@@ -0,0 +1,5 @@
+[style]
+based_on_style = pep8
+blank_line_before_nested_class_or_def = false
+spaces_before_comment = 2
+split_before_logical_operator = true

README.md

@@ -0,0 +1,11 @@
+---
+title: DeepDanbooru
+emoji: 🏃
+colorFrom: gray
+colorTo: purple
+sdk: gradio
+sdk_version: 3.34.0
+app_file: app.py
+pinned: false
+duplicated_from: hysts/DeepDanbooru
+---

app.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+
+import os
+import pathlib
+import tarfile
+
+import deepdanbooru as dd
+import gradio as gr
+import huggingface_hub
+import numpy as np
+import PIL.Image
+import tensorflow as tf
+
+DESCRIPTION = '# [KichangKim/DeepDanbooru](https://github.com/KichangKim/DeepDanbooru)'
+
+
+def load_sample_image_paths() -> list[pathlib.Path]:
+    image_dir = pathlib.Path('images')
+    if not image_dir.exists():
+        path = huggingface_hub.hf_hub_download(
+            'public-data/sample-images-TADNE',
+            'images.tar.gz',
+            repo_type='dataset')
+        with tarfile.open(path) as f:
+            f.extractall()
+    return sorted(image_dir.glob('*'))
+
+
+def load_model() -> tf.keras.Model:
+    path = huggingface_hub.hf_hub_download('public-data/DeepDanbooru',
+                                           'model-resnet_custom_v3.h5')
+    model = tf.keras.models.load_model(path)
+    return model
+
+
+def load_labels() -> list[str]:
+    path = huggingface_hub.hf_hub_download('public-data/DeepDanbooru',
+                                           'tags.txt')
+    with open(path) as f:
+        labels = [line.strip() for line in f.readlines()]
+    return labels
+
+
+model = load_model()
+labels = load_labels()
+
+
+def predict(image: PIL.Image.Image,
+            score_threshold: float) -> dict[str, float]:
+    _, height, width, _ = model.input_shape
+    image = np.asarray(image)
+    image = tf.image.resize(image,
+                            size=(height, width),
+                            method=tf.image.ResizeMethod.AREA,
+                            preserve_aspect_ratio=True)
+    image = image.numpy()
+    image = dd.image.transform_and_pad_image(image, width, height)
+    image = image / 255.
+    probs = model.predict(image[None, ...])[0]
+    probs = probs.astype(float)
+    res = dict()
+    for prob, label in zip(probs.tolist(), labels):
+        if prob < score_threshold:
+            continue
+        res[label] = prob
+    return res
+
+
+image_paths = load_sample_image_paths()
+examples = [[path.as_posix(), 0.5] for path in image_paths]
+
+with gr.Blocks(css='style.css') as demo:
+    gr.Markdown(DESCRIPTION)
+    with gr.Row():
+        with gr.Column():
+            image = gr.Image(label='Input', type='pil')
+            score_threshold = gr.Slider(label='Score threshold',
+                                        minimum=0,
+                                        maximum=1,
+                                        step=0.05,
+                                        value=0.5)
+            run_button = gr.Button('Run')
+        with gr.Column():
+            result = gr.Label(label='Output')
+    gr.Examples(examples=examples,
+                inputs=[image, score_threshold],
+                outputs=result,
+                fn=predict,
+                cache_examples=os.getenv('CACHE_EXAMPLES') == '1')
+
+    run_button.click(fn=predict,
+                     inputs=[image, score_threshold],
+                     outputs=result,
+                     api_name='predict')
+demo.queue().launch()

requirements.txt

@@ -0,0 +1,3 @@
+pillow>=9.0.0
+tensorflow>=2.7.0
+git+https://github.com/KichangKim/DeepDanbooru@v3-20200915-sgd-e30#egg=deepdanbooru

style.css

@@ -0,0 +1,3 @@
+h1 {
+  text-align: center;
+}