jianshang

.gitignore

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

.pre-commit-config.yaml

@@ -0,0 +1,35 @@
+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.10.1
+ hooks:
+ - id: isort
+- repo: https://github.com/pre-commit/mirrors-mypy
+ rev: v0.812
+ hooks:
+ - id: mypy
+ args: ['--ignore-missing-imports']
+- 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

@@ -4,6 +4,7 @@ emoji: 🏃
 colorFrom: gray
 colorTo: purple
 sdk: gradio
+sdk_version: 3.6
 app_file: app.py
 pinned: false
 ---

app.py

@@ -0,0 +1,122 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+
+import argparse
+import functools
+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
+
+TITLE = 'KichangKim/DeepDanbooru'
+DESCRIPTION = 'This is an unofficial demo for https://github.com/KichangKim/DeepDanbooru.'
+ARTICLE = '<center><img src="https://visitor-badge.glitch.me/badge?page_id=hysts.deepdanbooru" alt="visitor badge"/></center>'
+
+HF_TOKEN = os.environ['HF_TOKEN']
+MODEL_REPO = 'hysts/DeepDanbooru'
+MODEL_FILENAME = 'model-resnet_custom_v3.h5'
+LABEL_FILENAME = 'tags.txt'
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--score-slider-step', type=float, default=0.05)
+ parser.add_argument('--score-threshold', type=float, default=0.5)
+ parser.add_argument('--share', action='store_true')
+ return parser.parse_args()
+
+
+def load_sample_image_paths() -> list[pathlib.Path]:
+ image_dir = pathlib.Path('images')
+ if not image_dir.exists():
+ dataset_repo = 'hysts/sample-images-TADNE'
+ path = huggingface_hub.hf_hub_download(dataset_repo,
+ 'images.tar.gz',
+ repo_type='dataset',
+ use_auth_token=HF_TOKEN)
+ 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(MODEL_REPO,
+ MODEL_FILENAME,
+ use_auth_token=HF_TOKEN)
+ model = tf.keras.models.load_model(path)
+ return model
+
+
+def load_labels() -> list[str]:
+ path = huggingface_hub.hf_hub_download(MODEL_REPO,
+ LABEL_FILENAME,
+ use_auth_token=HF_TOKEN)
+ with open(path) as f:
+ labels = [line.strip() for line in f.readlines()]
+ return labels
+
+
+def predict(image: PIL.Image.Image, score_threshold: float,
+ model: tf.keras.Model, labels: list[str]) -> 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
+
+
+def main():
+ args = parse_args()
+
+ image_paths = load_sample_image_paths()
+ examples = [[path.as_posix(), args.score_threshold]
+ for path in image_paths]
+
+ model = load_model()
+ labels = load_labels()
+
+ func = functools.partial(predict, model=model, labels=labels)
+
+ gr.Interface(
+ func,
+ [
+ gr.Image(type='pil', label='Input'),
+ gr.Slider(0,
+ 1,
+ step=args.score_slider_step,
+ value=args.score_threshold,
+ label='Score Threshold'),
+ ],
+ gr.Label(label='Output'),
+ examples=examples,
+ title=TITLE,
+ description=DESCRIPTION,
+ article=ARTICLE,
+ allow_flagging='never',
+ ).launch(
+ enable_queue=True,
+ share=args.share,
+ )
+
+
+if __name__ == '__main__':
+ main()

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