reslove

DeepDanbooru/.gitignore

@@ -0,0 +1,15 @@
+.vs/
+.vscode/
+__pycache__/
+/test*
+test.py
+/logs
+/uploads
+*.jpg
+*.jpeg
+*.png
+*.bmp
+*.gif
+*.h5
+*.txt
+*.json

DeepDanbooru/.travis.yml

@@ -0,0 +1,12 @@
+language: python
+python:
+  - "3.6"
+# command to install dependencies
+install:
+  - pip install -e .[test]
+  - pip install tensorflow
+# command to run tests
+script:
+  - pytest
+  - flake8
+  - mypy .

DeepDanbooru/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 Kichang Kim
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

DeepDanbooru/README.md

@@ -0,0 +1,99 @@
+# DeepDanbooru
+[![Python](https://img.shields.io/badge/python-3.6-green)](https://www.python.org/doc/versions/)
+[![GitHub](https://img.shields.io/github/license/KichangKim/DeepDanbooru)](https://opensource.org/licenses/MIT)
+[![Web](https://img.shields.io/badge/web%20demo-20191108-brightgreen)](http://kanotype.iptime.org:8003/deepdanbooru/)
+
+**DeepDanbooru** is anime-style girl image tag estimation system. You can estimate your images on my live demo site, [DeepDanbooru Web](http://kanotype.iptime.org:8003/deepdanbooru/).
+
+## Requirements
+DeepDanbooru is written by Python 3.6. Following packages are need to be installed.
+- tensorflow>=2.1.0
+- Click>=7.0
+- numpy>=1.16.2
+- requests>=2.22.0
+- scikit-image>=0.15.0
+- six>=1.13.0
+
+Or just use `requirements.txt`.
+```
+> pip install -r requirements.txt
+```
+
+alternatively you can install it with pip. Note that by default, tensorflow is not included.
+
+To install it with tensorflow, add `tensorflow` extra package.
+
+```
+> # default installation
+> pip install .
+> # with tensorflow package
+> pip install .[tensorflow]
+```
+
+
+## Usage
+1. Prepare dataset. If you don't have, you can use [DanbooruDownloader](https://github.com/KichangKim/DanbooruDownloader) for download the dataset of [Danbooru](https://danbooru.donmai.us/). If you want to make your own dataset, see [Dataset Structure](#dataset-structure) section.
+2. Create training project folder.
+```
+> deepdanbooru create-project [your_project_folder]
+```
+3. Prepare tag list. If you want to use latest tags, use following command. It downloads tag from Danbooru server.
+```
+> deepdanbooru download-tags [your_project_folder]
+```
+4. (Option) Filtering dataset. If you want to train with optional tags (rating and score), you should convert it as system tags.
+```
+> deepdanbooru make-training-database [your_dataset_sqlite_path] [your_filtered_sqlite_path]
+```
+5. Modify `project.json` in the project folder. You should change `database_path` setting to your actual sqlite file path.
+6. Start training.
+```
+> deepdanbooru train-project [your_project_folder]
+```
+7. Enjoy it.
+```
+> deepdanbooru evaluate [image_file_path or folder]... --project-path [your_project_folder] --allow-folder
+```
+
+## Dataset Structure
+DeepDanbooru uses following folder structure for input dataset. SQLite file can be any name, but must be located in same folder to `images` folder.
+```
+MyDataset/
+├── images/
+│   ├── 00/
+│   │   ├── 00000000000000000000000000000000.jpg
+│   │   ├── ...
+│   ├── 01/
+│   │   ├── ...
+│   └── ff/
+│       ├── ...
+└── my-dataset.sqlite
+```
+The core is SQLite database file. That file must be contains following table structure.
+```
+posts
+├── id (INTEGER)
+├── md5 (TEXT)
+├── file_ext (TEXT)
+├── tag_string (TEXT)
+└── tag_count_general (INTEGER)
+```
+The filename of image must be `[md5].[file_ext]`. If you use your own images, `md5` don't have to be actual MD5 hash value.
+
+`tag_string` is space splitted tag list, like `1girl ahoge long_hair`.
+
+`tag_count_general` is used for the project setting, `minimum_tag_count`. Images which has equal or larger value of `tag_count_general` are used for training.
+
+## Project Structure
+**Project** is minimal unit for training on DeepDanbooru. You can modify various parameters for training.
+```
+MyProject/
+├── project.json
+└── tags.txt
+```
+`tags.txt` contains all tags for estimating. You can make your own list or download latest tags from Danbooru server. It is simple newline-separated file like this:
+```
+1girl
+ahoge
+...
+```

DeepDanbooru/deepdanbooru/__init__.py

@@ -0,0 +1,8 @@
+import deepdanbooru.commands
+import deepdanbooru.data
+import deepdanbooru.extra
+import deepdanbooru.image
+import deepdanbooru.io
+import deepdanbooru.model
+import deepdanbooru.project
+import deepdanbooru.train

DeepDanbooru/deepdanbooru/__main__.py

@@ -0,0 +1,89 @@
+import sys
+
+import click
+
+import deepdanbooru as dd
+
+__version__ = '1.0.0'
+
+
+@click.version_option(prog_name='DeepDanbooru', version=__version__)
+@click.group()
+def main():
+    '''
+    AI based multi-label girl image classification system, implemented by using TensorFlow.
+    '''
+    pass
+
+
+@main.command('create-project')
+@click.argument('project_path', type=click.Path(exists=False, resolve_path=True, file_okay=False, dir_okay=True))
+def create_project(project_path):
+    dd.commands.create_project(project_path)
+
+
+@main.command('download-tags')
+@click.option('--limit', default=10000, help='Limit for each category tag count.')
+@click.option('--minimum-post-count', default=500, help='Minimum post count for tag.')
+@click.option('--overwrite', help='Overwrite tags if exists.', is_flag=True)
+@click.argument('path', type=click.Path(exists=False, resolve_path=True, file_okay=False, dir_okay=True))
+def download_tags(path, limit, minimum_post_count, overwrite):
+    dd.commands.download_tags(path, limit, minimum_post_count, overwrite)
+
+
+@main.command('make-training-database')
+@click.argument('source_path', type=click.Path(exists=True, resolve_path=True, file_okay=True, dir_okay=False), nargs=1, required=True)
+@click.argument('output_path', type=click.Path(exists=False, resolve_path=True, file_okay=True, dir_okay=False), nargs=1, required=True)
+@click.option('--start-id', default=1, help='Start id.', )
+@click.option('--end-id', default=sys.maxsize, help='End id.')
+@click.option('--use-deleted', help='Use deleted posts.', is_flag=True)
+@click.option('--chunk-size', default=5000000, help='Chunk size for internal processing.')
+@click.option('--overwrite', help='Overwrite tags if exists.', is_flag=True)
+@click.option('--vacuum', help='Execute VACUUM command after making database.', is_flag=True)
+def make_training_database(source_path, output_path, start_id, end_id, use_deleted, chunk_size, overwrite, vacuum):
+    dd.commands.make_training_database(source_path, output_path, start_id, end_id,
+                                       use_deleted, chunk_size, overwrite, vacuum)
+
+
+@main.command('train-project')
+@click.argument('project_path', type=click.Path(exists=True, resolve_path=True, file_okay=False, dir_okay=True))
+@click.option('--source-model', type=click.Path(exists=True, resolve_path=True, file_okay=True, dir_okay=False))
+def train_project(project_path, source_model):
+    dd.commands.train_project(project_path, source_model)
+
+
+@main.command('evaluate-project', help='Evaluate the project. If the target path is folder, it evaulates all images recursively.')
+@click.argument('project_path', type=click.Path(exists=True, resolve_path=True, file_okay=False, dir_okay=True))
+@click.argument('target_path', type=click.Path(exists=True, resolve_path=True, file_okay=True, dir_okay=True))
+@click.option('--threshold', help='Threshold for tag estimation.', default=0.5)
+def evaluate_project(project_path, target_path, threshold):
+    dd.commands.evaluate_project(project_path, target_path, threshold)
+
+
+@main.command('grad-cam', help='Experimental feature. Calculate activation map using Grad-CAM.')
+@click.argument('project_path', type=click.Path(exists=True, resolve_path=True, file_okay=False, dir_okay=True))
+@click.argument('target_path', type=click.Path(exists=True, resolve_path=True, file_okay=True, dir_okay=True))
+@click.argument('output_path', type=click.Path(resolve_path=True, file_okay=False, dir_okay=True), default='.')
+@click.option('--threshold', help='Threshold for tag estimation.', default=0.5)
+def grad_cam(project_path, target_path, output_path, threshold):
+    dd.commands.grad_cam(project_path, target_path, output_path, threshold)
+
+
+@main.command('evaluate', help='Evaluate model by estimating image tag.')
+@click.argument('target_paths', nargs=-1, type=click.Path(exists=True, resolve_path=True, file_okay=True, dir_okay=True))
+@click.option('--project-path', type=click.Path(exists=True, resolve_path=True, file_okay=False, dir_okay=True),
+              help='Project path. If you want to use specific model and tags, use --model-path and --tags-path options.')
+@click.option('--model-path', type=click.Path(exists=True, resolve_path=True, file_okay=True, dir_okay=False))
+@click.option('--tags-path', type=click.Path(exists=True, resolve_path=True, file_okay=True, dir_okay=False))
+@click.option('--threshold', default=0.5)
+@click.option('--allow-gpu', default=False, is_flag=True)
+@click.option('--compile/--no-compile', 'compile_model', default=False)
+@click.option('--allow-folder', default=False, is_flag=True, help='If this option is enabled, TARGET_PATHS can be folder path and all images (using --folder-filters) in that folder is estimated recursively. If there are file and folder which has same name, the file is skipped and only folder is used.')
+@click.option('--folder-filters', default='*.[Pp][Nn][Gg],*.[Jj][Pp][Gg],*.[Jj][Pp][Ee][Gg],*.[Gg][Ii][Ff]', help='Glob pattern for searching image files in folder. You can specify multiple patterns by separating comma. This is used when --allow-folder is enabled. Default:*.[Pp][Nn][Gg],*.[Jj][Pp][Gg],*.[Jj][Pp][Ee][Gg],*.[Gg][Ii][Ff]')
+@click.option('--verbose', default=False, is_flag=True)
+def evaluate(target_paths, project_path, model_path, tags_path, threshold, allow_gpu, compile_model, allow_folder, folder_filters, verbose):
+    dd.commands.evaluate(target_paths, project_path, model_path, tags_path, threshold, allow_gpu, compile_model, allow_folder, folder_filters, verbose)
+
+
+if __name__ == '__main__':
+    main()

DeepDanbooru/deepdanbooru/commands/__init__.py

@@ -0,0 +1,7 @@
+from .create_project import create_project
+from .download_tags import download_tags
+from .make_training_database import make_training_database
+from .train_project import train_project
+from .evaluate_project import evaluate_project
+from .grad_cam import grad_cam
+from .evaluate import evaluate, evaluate_image

DeepDanbooru/deepdanbooru/commands/create_project.py

@@ -0,0 +1,15 @@
+import os
+
+import deepdanbooru as dd
+
+
+def create_project(project_path):
+    """
+    Create new project with default parameters.
+    """
+    dd.io.try_create_directory(project_path)
+    project_context_path = os.path.join(project_path, 'project.json')
+    dd.io.serialize_as_json(
+        dd.project.DEFAULT_PROJECT_CONTEXT, project_context_path)
+
+    print(f'New project was successfully created. ({project_path})')

DeepDanbooru/deepdanbooru/commands/download_tags.py

@@ -0,0 +1,162 @@
+import os
+import time
+
+import requests
+
+import deepdanbooru as dd
+
+
+def download_category_tags(category, minimum_post_count, limit, page_size=1000, order='count'):
+    category_to_index = {
+        'general': 0,
+        'artist': 1,
+        'copyright': 3,
+        'character': 4
+    }
+
+    gold_only_tags = ['loli', 'shota', 'toddlercon']
+
+    if category not in category_to_index:
+        raise Exception(f'Not supported category : {category}')
+
+    category_index = category_to_index[category]
+
+    parameters = {
+        'limit': page_size,
+        'page': 1,
+        'search[order]': order,
+        'search[category]': category_index
+    }
+
+    request_url = 'https://danbooru.donmai.us/tags.json'
+
+    tags = set()
+
+    while True:
+        response = requests.get(request_url, params=parameters)
+        response_json = response.json()
+
+        response_tags = [tag_json['name']
+                         for tag_json in response_json if tag_json['post_count'] >= minimum_post_count]
+
+        if not response_tags:
+            break
+
+        is_full = False
+
+        for tag in response_tags:
+            if tag in gold_only_tags:
+                continue
+
+            tags.add(tag)
+
+            if len(tags) >= limit:
+                is_full = True
+                break
+
+        if is_full:
+            break
+        else:
+            parameters['page'] += 1
+
+    return tags
+
+
+def download_tags(project_path, limit, minimum_post_count, is_overwrite):
+    print(
+        f'Start downloading tags ... (limit:{limit}, minimum_post_count:{minimum_post_count})')
+
+    log = {
+        'date': time.strftime("%Y/%m/%d %H:%M:%S"),
+        'limit': limit,
+        'minimum_post_count': minimum_post_count
+    }
+
+    system_tags = [
+        'rating:safe',
+        'rating:questionable',
+        'rating:explicit',
+        # 'score:very_bad',
+        # 'score:bad',
+        # 'score:average',
+        # 'score:good',
+        # 'score:very_good',
+    ]
+
+    category_definitions = [
+        {
+            'category_name': 'General',
+            'category': 'general',
+            'path': os.path.join(project_path, 'tags-general.txt'),
+        },
+        # {
+        #    'category_name': 'Artist',
+        #    'category': 'artist',
+        #    'path': os.path.join(path, 'tags-artist.txt'),
+        # },
+        # {
+        #    'category_name': 'Copyright',
+        #    'category': 'copyright',
+        #    'path': os.path.join(path, 'tags-copyright.txt'),
+        # },
+        {
+            'category_name': 'Character',
+            'category': 'character',
+            'path': os.path.join(project_path, 'tags-character.txt'),
+        },
+    ]
+
+    all_tags_path = os.path.join(project_path, 'tags.txt')
+
+    if not is_overwrite and os.path.exists(all_tags_path):
+        raise Exception(f'Tags file is already exists : {all_tags_path}')
+
+    dd.io.try_create_directory(os.path.dirname(all_tags_path))
+    dd.io.serialize_as_json(
+        log, os.path.join(project_path, 'tags_log.json'))
+
+    categories_for_web = []
+    categories_for_web_path = os.path.join(project_path, 'categories.json')
+    tag_start_index = 0
+
+    total_tags_count = 0
+
+    with open(all_tags_path, 'w') as all_tags_stream:
+        for category_definition in category_definitions:
+            category = category_definition['category']
+            category_tags_path = category_definition['path']
+
+            print(f'{category} tags are downloading ...')
+            tags = download_category_tags(category, minimum_post_count, limit)
+
+            tags = dd.extra.natural_sorted(tags)
+            tag_count = len(tags)
+            if tag_count == 0:
+                print(f'{category} tags are not exists.')
+                continue
+            else:
+                print(f'{tag_count} tags are downloaded.')
+
+            with open(category_tags_path, 'w') as category_tags_stream:
+                for tag in tags:
+                    category_tags_stream.write(f'{tag}\n')
+                    all_tags_stream.write(f'{tag}\n')
+
+            categories_for_web.append(
+                {'name': category_definition['category_name'], 'start_index': tag_start_index})
+
+            tag_start_index += len(tags)
+            total_tags_count += tag_count
+
+        for tag in system_tags:
+            all_tags_stream.write(f'{tag}\n')
+
+        categories_for_web.append(
+            {'name': 'System', 'start_index': total_tags_count}
+        )
+
+    dd.io.serialize_as_json(categories_for_web, categories_for_web_path)
+
+    print(f'Total {total_tags_count} tags are downloaded.')
+
+    print('All processes are complete.')

DeepDanbooru/deepdanbooru/commands/evaluate.py

@@ -0,0 +1,77 @@
+import os
+from typing import Any, Iterable, List, Tuple, Union
+
+import six
+import tensorflow as tf
+
+import deepdanbooru as dd
+
+
+def evaluate_image(
+    image_input: Union[str, six.BytesIO], model: Any, tags: List[str], threshold: float
+) -> Iterable[Tuple[str, float]]:
+    width = model.input_shape[2]
+    height = model.input_shape[1]
+
+    image = dd.data.load_image_for_evaluate(
+        image_input, width=width, height=height)
+
+    image_shape = image.shape
+    image = image.reshape(
+        (1, image_shape[0], image_shape[1], image_shape[2]))
+    y = model.predict(image)[0]
+
+    result_dict = {}
+
+    for i, tag in enumerate(tags):
+        result_dict[tag] = y[i]
+
+    for tag in tags:
+        if result_dict[tag] >= threshold:
+            yield tag, result_dict[tag]
+
+
+def evaluate(target_paths, project_path, model_path, tags_path, threshold, allow_gpu, compile_model, allow_folder, folder_filters, verbose):
+    if not allow_gpu:
+        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+
+    if not model_path and not project_path:
+        raise Exception('You must provide project path or model path.')
+
+    if not tags_path and not project_path:
+        raise Exception('You must provide project path or tags path.')
+
+    target_image_paths = []
+
+    for target_path in target_paths:
+        if allow_folder and not os.path.isfile(target_path):
+            target_image_paths.extend(dd.io.get_image_file_paths_recursive(target_path, folder_filters))
+        else:
+            target_image_paths.append(target_path)
+
+    target_image_paths = dd.extra.natural_sorted(target_image_paths)
+
+    if model_path:
+        if verbose:
+            print(f'Loading model from {model_path} ...')
+        model = tf.keras.models.load_model(model_path, compile=compile_model)
+    else:
+        if verbose:
+            print(f'Loading model from project {project_path} ...')
+        model = dd.project.load_model_from_project(project_path, compile_model=compile_model)
+
+    if tags_path:
+        if verbose:
+            print(f'Loading tags from {tags_path} ...')
+        tags = dd.data.load_tags(tags_path)
+    else:
+        if verbose:
+            print(f'Loading tags from project {project_path} ...')
+        tags = dd.project.load_tags_from_project(project_path)
+
+    for image_path in target_image_paths:
+        print(f'Tags of {image_path}:')
+        for tag, score in evaluate_image(image_path, model, tags, threshold):
+            print(f'({score:05.3f}) {tag}')
+
+        print()

DeepDanbooru/deepdanbooru/commands/evaluate_project.py

@@ -0,0 +1,51 @@
+import os
+
+import deepdanbooru as dd
+
+
+def evaluate_project(project_path, target_path, threshold):
+    if not os.path.exists(target_path):
+        raise Exception(f'Target path {target_path} is not exists.')
+
+    if os.path.isfile(target_path):
+        taget_image_paths = [target_path]
+
+    else:
+        patterns = [
+            '*.[Pp][Nn][Gg]',
+            '*.[Jj][Pp][Gg]',
+            '*.[Jj][Pp][Ee][Gg]',
+            '*.[Gg][Ii][Ff]'
+        ]
+
+        taget_image_paths = dd.io.get_file_paths_in_directory(
+            target_path, patterns)
+
+        taget_image_paths = dd.extra.natural_sorted(taget_image_paths)
+
+    project_context, model, tags = dd.project.load_project(project_path)
+
+    width = project_context['image_width']
+    height = project_context['image_height']
+
+    for image_path in taget_image_paths:
+        image = dd.data.load_image_for_evaluate(
+            image_path, width=width, height=height)
+
+        image_shape = image.shape
+        # image = image.astype(np.float16)
+        image = image.reshape(
+            (1, image_shape[0], image_shape[1], image_shape[2]))
+        y = model.predict(image)[0]
+
+        result_dict = {}
+
+        for i, tag in enumerate(tags):
+            result_dict[tag] = y[i]
+
+        print(f'Tags of {image_path}:')
+        for tag in tags:
+            if result_dict[tag] >= threshold:
+                print(f'({result_dict[tag]:05.3f}) {tag}')
+
+        print()

DeepDanbooru/deepdanbooru/commands/grad_cam.py

@@ -0,0 +1,111 @@
+import os
+
+import tensorflow as tf
+import numpy as np
+from PIL import Image
+import deepdanbooru as dd
+from scipy import ndimage
+
+
+@tf.function
+def get_gradient(model, x, output_mask):
+    with tf.GradientTape() as tape:
+        output = model(x)
+        gradcam_loss = tf.reduce_sum(tf.multiply(output_mask, output))
+
+    return tape.gradient(gradcam_loss, x)
+
+
+def norm_clip_grads(grads):
+    upper_quantile = np.quantile(grads, 0.99)
+    lower_quantile = np.quantile(grads, 0.01)
+    clipped_grads = np.abs(np.clip(grads, lower_quantile, upper_quantile))
+
+    return clipped_grads / np.max(clipped_grads)
+
+
+def filter_grads(grads):
+    return ndimage.median_filter(grads, 10)
+
+
+def to_onehot(length, index):
+    value = np.zeros(shape=(1, length), dtype=np.float32)
+    value[0, index] = 1.0
+    return value
+
+
+def grad_cam(project_path, target_path, output_path, threshold):
+    # os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+
+    if not os.path.exists(target_path):
+        raise Exception(f'Target path {target_path} is not exists.')
+
+    if os.path.isfile(target_path):
+        taget_image_paths = [target_path]
+    else:
+        patterns = [
+            '*.[Pp][Nn][Gg]',
+            '*.[Jj][Pp][Gg]',
+            '*.[Jj][Pp][Ee][Gg]',
+            '*.[Gg][Ii][Ff]'
+        ]
+
+        taget_image_paths = dd.io.get_file_paths_in_directory(
+            target_path, patterns)
+
+        taget_image_paths = dd.extra.natural_sorted(taget_image_paths)
+
+    model = dd.project.load_model_from_project(project_path)
+    tags = dd.project.load_tags_from_project(project_path)
+    width = model.input_shape[2]
+    height = model.input_shape[1]
+
+    dd.io.try_create_directory(output_path)
+
+    for image_path in taget_image_paths:
+        image = dd.data.load_image_for_evaluate(
+            image_path, width=width, height=height)
+        image_name = os.path.splitext(os.path.basename(image_path))[0]
+
+        image_folder = os.path.join(output_path, image_name)
+        dd.io.try_create_directory(image_folder)
+
+        Image.fromarray(np.uint8(
+            image * 255.0)).save(os.path.join(image_folder, f'input.png'))
+        image_for_result = image
+        image_shape = image.shape
+        y = model.predict(image.reshape(
+            (1, image_shape[0], image_shape[1], image_shape[2])))[0]
+
+        result_dict = {}
+
+        estimated_tags = []
+
+        for i, tag in enumerate(tags):
+            result_dict[tag] = y[i]
+
+            if y[i] >= threshold:
+                estimated_tags.append((i, tag))
+
+        print(f'Tags of {image_path}:')
+
+        for tag in tags:
+            if result_dict[tag] >= threshold:
+                print(f'({result_dict[tag]:05.3f}) {tag}')
+
+        image = image.astype(np.float32)
+
+        for estimated_tag in estimated_tags:
+            print(f'Calculating grad-cam ... ({estimated_tag[1]})')
+            grads = get_gradient(model, tf.Variable(
+                [image]), to_onehot(len(tags), estimated_tag[0]))[0]
+            print('Normalizing gradients ...')
+            grads = norm_clip_grads(grads)
+            print('Filtering gradients ...')
+            grads = filter_grads(grads)
+            Image.fromarray(np.uint8(grads * 255.0)
+                            ).save(os.path.join(image_folder, f'result-{estimated_tag[1]}.png'.replace(':', '_').replace('/', '_')))
+            mask_array = np.stack([np.max(
+                grads, axis=-1)]*3, axis=2)
+            Image.fromarray(np.uint8(np.multiply(image_for_result, mask_array)
+                                     * 255.0)).save(os.path.join(image_folder, f'result-{estimated_tag[1]}-masked.png'.replace(':', '_').replace('/', '_')))

DeepDanbooru/deepdanbooru/commands/make_training_database.py

@@ -0,0 +1,122 @@
+import os
+import sqlite3
+
+
+def make_training_database(source_path, output_path, start_id, end_id,
+                           use_deleted, chunk_size, overwrite, vacuum):
+    '''
+    Make sqlite database for training. Also add system tags.
+    '''
+    if source_path == output_path:
+        raise Exception('Source path and output path is equal.')
+
+    if os.path.exists(output_path):
+        if overwrite:
+            os.remove(output_path)
+        else:
+            raise Exception(f'{output_path} is already exists.')
+
+    source_connection = sqlite3.connect(source_path)
+    source_connection.row_factory = sqlite3.Row
+    source_cursor = source_connection.cursor()
+
+    output_connection = sqlite3.connect(output_path)
+    output_connection.row_factory = sqlite3.Row
+    output_cursor = output_connection.cursor()
+
+    table_name = 'posts'
+    id_column_name = 'id'
+    md5_column_name = 'md5'
+    extension_column_name = 'file_ext'
+    tags_column_name = 'tag_string'
+    tag_count_general_column_name = 'tag_count_general'
+    rating_column_name = 'rating'
+    score_column_name = 'score'
+    deleted_column_name = 'is_deleted'
+
+    # Create output table
+    print('Creating table ...')
+    output_cursor.execute(f"""CREATE TABLE {table_name} (
+        {id_column_name} INTEGER NOT NULL PRIMARY KEY,
+        {md5_column_name} TEXT,
+        {extension_column_name} TEXT,
+        {tags_column_name} TEXT,
+        {tag_count_general_column_name} INTEGER )""")
+    output_connection.commit()
+    print('Creating table is complete.')
+
+    current_start_id = start_id
+
+    while True:
+        print(
+            f'Fetching source rows ... ({current_start_id}~)')
+        source_cursor.execute(
+            f"""SELECT
+                {id_column_name},{md5_column_name},{extension_column_name},{tags_column_name},{tag_count_general_column_name},{rating_column_name},{score_column_name},{deleted_column_name}
+            FROM {table_name} WHERE ({id_column_name} >= ?) ORDER BY {id_column_name} ASC LIMIT ?""",
+            (current_start_id, chunk_size))
+
+        rows = source_cursor.fetchall()
+
+        if not rows:
+            break
+
+        insert_params = []
+
+        for row in rows:
+            post_id = row[id_column_name]
+            md5 = row[md5_column_name]
+            extension = row[extension_column_name]
+            tags = row[tags_column_name]
+            general_tag_count = row[tag_count_general_column_name]
+            rating = row[rating_column_name]
+            # score = row[score_column_name]
+            is_deleted = row[deleted_column_name]
+
+            if post_id > end_id:
+                break
+
+            if is_deleted and not use_deleted:
+                continue
+
+            if rating == 's':
+                tags += f' rating:safe'
+            elif rating == 'q':
+                tags += f' rating:questionable'
+            elif rating == 'e':
+                tags += f' rating:explicit'
+
+            # if score < -6:
+            #     tags += f' score:very_bad'
+            # elif score >= -6 and score < 0:
+            #     tags += f' score:bad'
+            # elif score >= 0 and score < 7:
+            #     tags += f' score:average'
+            # elif score >= 7 and score < 13:
+            #     tags += f' score:good'
+            # elif score >= 13:
+            #     tags += f' score:very_good'
+
+            insert_params.append(
+                (post_id, md5, extension, tags, general_tag_count))
+
+        if insert_params:
+            print('Inserting ...')
+            output_cursor.executemany(
+                f"""INSERT INTO {table_name} (
+                {id_column_name},{md5_column_name},{extension_column_name},{tags_column_name},{tag_count_general_column_name})
+                values (?, ?, ?, ?, ?)""", insert_params)
+            output_connection.commit()
+
+        current_start_id = rows[-1][id_column_name] + 1
+
+        if current_start_id > end_id or len(rows) < chunk_size:
+            break
+
+    if vacuum:
+        print('Vacuum ...')
+        output_cursor.execute('vacuum')
+        output_connection.commit()
+
+    source_connection.close()
+    output_connection.close()

DeepDanbooru/deepdanbooru/commands/train_project.py

@@ -0,0 +1,222 @@
+import os
+import random
+import time
+import datetime
+
+import tensorflow as tf
+
+import deepdanbooru as dd
+
+
+def train_project(project_path, source_model):
+    project_context_path = os.path.join(project_path, 'project.json')
+    project_context = dd.io.deserialize_from_json(project_context_path)
+
+    width = project_context['image_width']
+    height = project_context['image_height']
+    database_path = project_context['database_path']
+    minimum_tag_count = project_context['minimum_tag_count']
+    model_type = project_context['model']
+    optimizer_type = project_context['optimizer']
+    learning_rate = project_context['learning_rate'] if 'learning_rate' in project_context else 0.001
+    learning_rates = project_context['learning_rates'] if 'learning_rates' in project_context else None
+    minibatch_size = project_context['minibatch_size']
+    epoch_count = project_context['epoch_count']
+    export_model_per_epoch = project_context[
+        'export_model_per_epoch'] if 'export_model_per_epoch' in project_context else 10
+    checkpoint_frequency_mb = project_context['checkpoint_frequency_mb']
+    console_logging_frequency_mb = project_context['console_logging_frequency_mb']
+    rotation_range = project_context['rotation_range']
+    scale_range = project_context['scale_range']
+    shift_range = project_context['shift_range']
+
+    # disable PNG warning
+    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
+    # tf.logging.set_verbosity(tf.logging.ERROR)
+
+    # tf.keras.backend.set_epsilon(1e-6)
+    # tf.keras.mixed_precision.experimental.set_policy('infer_float32_vars')
+    # tf.config.gpu.set_per_process_memory_growth(True)
+
+    if optimizer_type == 'adam':
+        optimizer = tf.optimizers.Adam(learning_rate)
+        print('Using Adam optimizer ... ')
+    elif optimizer_type == 'sgd':
+        optimizer = tf.optimizers.SGD(
+            learning_rate, momentum=0.9, nesterov=True)
+        print('Using SGD optimizer ... ')
+    elif optimizer_type == 'rmsprop':
+        optimizer = tf.optimizers.RMSprop(learning_rate)
+        print('Using RMSprop optimizer ... ')
+    else:
+        raise Exception(
+            f"Not supported optimizer : {optimizer_type}")
+
+    if model_type == 'resnet_152':
+        model_delegate = dd.model.resnet.create_resnet_152
+    elif model_type == 'resnet_custom_v1':
+        model_delegate = dd.model.resnet.create_resnet_custom_v1
+    elif model_type == 'resnet_custom_v2':
+        model_delegate = dd.model.resnet.create_resnet_custom_v2
+    elif model_type == 'resnet_custom_v3':
+        model_delegate = dd.model.resnet.create_resnet_custom_v3
+    elif model_type == 'resnet_custom_v4':
+        model_delegate = dd.model.resnet.create_resnet_custom_v4
+    else:
+        raise Exception(f'Not supported model : {model_type}')
+
+    print('Loading tags ... ')
+    tags = dd.project.load_tags_from_project(project_path)
+    output_dim = len(tags)
+
+    print(f'Creating model ({model_type}) ... ')
+    # tf.keras.backend.set_learning_phase(1)
+
+    if source_model:
+        model = tf.keras.models.load_model(source_model)
+        print(f'Model : {model.input_shape} -> {model.output_shape} (loaded from {source_model})')
+    else:
+        inputs = tf.keras.Input(shape=(height, width, 3),
+                                dtype=tf.float32)  # HWC
+        ouputs = model_delegate(inputs, output_dim)
+        model = tf.keras.Model(inputs=inputs, outputs=ouputs, name=model_type)
+        print(f'Model : {model.input_shape} -> {model.output_shape}')
+
+    model.compile(optimizer=optimizer, loss=tf.keras.losses.BinaryCrossentropy(),
+                  metrics=[tf.keras.metrics.Precision(), tf.keras.metrics.Recall()])
+
+    print(f'Loading database ... ')
+    image_records = dd.data.load_image_records(
+        database_path, minimum_tag_count)
+
+    # Checkpoint variables
+    used_epoch = tf.Variable(0, dtype=tf.int64)
+    used_minibatch = tf.Variable(0, dtype=tf.int64)
+    used_sample = tf.Variable(0, dtype=tf.int64)
+    offset = tf.Variable(0, dtype=tf.int64)
+    random_seed = tf.Variable(0, dtype=tf.int64)
+
+    checkpoint = tf.train.Checkpoint(
+        optimizer=optimizer,
+        model=model,
+        used_epoch=used_epoch,
+        used_minibatch=used_minibatch,
+        used_sample=used_sample,
+        offset=offset,
+        random_seed=random_seed)
+
+    manager = tf.train.CheckpointManager(
+        checkpoint=checkpoint,
+        directory=os.path.join(project_path, 'checkpoints'),
+        max_to_keep=3)
+
+    if manager.latest_checkpoint:
+        print(f"Checkpoint exists. Continuing training ... ({datetime.datetime.now()})")
+        checkpoint.restore(manager.latest_checkpoint)
+        print(f'used_epoch={int(used_epoch)}, used_minibatch={int(used_minibatch)}, used_sample={int(used_sample)}, offset={int(offset)}, random_seed={int(random_seed)}')
+    else:
+        print(f'No checkpoint. Starting new training ... ({datetime.datetime.now()})')
+
+    epoch_size = len(image_records)
+    slice_size = minibatch_size * checkpoint_frequency_mb
+    loss_sum = 0.0
+    loss_count = 0
+    used_sample_sum = 0
+    last_time = time.time()
+
+    while int(used_epoch) < epoch_count:
+        print(f'Shuffling samples (epoch {int(used_epoch)}) ... ')
+        epoch_random = random.Random(int(random_seed))
+        epoch_random.shuffle(image_records)
+
+        # Udpate learning rate
+        if learning_rates:
+            for learning_rate_per_epoch in learning_rates:
+                if learning_rate_per_epoch['used_epoch'] <= int(used_epoch):
+                    learning_rate = learning_rate_per_epoch['learning_rate']
+        print(f'Trying to change learning rate to {learning_rate} ...')
+        optimizer.learning_rate.assign(learning_rate)
+        print(f'Learning rate is changed to {optimizer.learning_rate} ...')
+
+        while int(offset) < epoch_size:
+            image_records_slice = image_records[int(offset):min(
+                int(offset) + slice_size, epoch_size)]
+
+            image_paths = [image_record[0]
+                           for image_record in image_records_slice]
+            tag_strings = [image_record[1]
+                           for image_record in image_records_slice]
+
+            dataset_wrapper = dd.data.DatasetWrapper(
+                (image_paths, tag_strings), tags, width, height, scale_range=scale_range, rotation_range=rotation_range, shift_range=shift_range)
+            dataset = dataset_wrapper.get_dataset(minibatch_size)
+
+            for (x_train, y_train) in dataset:
+                sample_count = x_train.shape[0]
+
+                step_result = model.train_on_batch(
+                    x_train, y_train, reset_metrics=False)
+
+                used_minibatch.assign_add(1)
+                used_sample.assign_add(sample_count)
+                used_sample_sum += sample_count
+                loss_sum += step_result[0]
+                loss_count += 1
+
+                if int(used_minibatch) % console_logging_frequency_mb == 0:
+                    # calculate logging informations
+                    current_time = time.time()
+                    delta_time = current_time - last_time
+                    step_metric_precision = step_result[1]
+                    step_metric_recall = step_result[2]
+                    if step_metric_precision + step_metric_recall > 0.0:
+                        step_metric_f1_score = 2.0 * \
+                            (step_metric_precision * step_metric_recall) / \
+                            (step_metric_precision + step_metric_recall)
+                    else:
+                        step_metric_f1_score = 0.0
+                    average_loss = loss_sum / float(loss_count)
+                    samples_per_seconds = float(
+                        used_sample_sum) / max(delta_time, 0.001)
+                    progress = float(int(used_sample)) / \
+                        float(epoch_size * epoch_count) * 100.0
+                    remain_seconds = float(
+                        epoch_size * epoch_count - int(used_sample)) / max(samples_per_seconds, 0.001)
+                    eta_datetime = datetime.datetime.now() + datetime.timedelta(seconds=remain_seconds)
+                    eta_datetime_string = eta_datetime.strftime(
+                        '%Y-%m-%d %H:%M:%S')
+                    print(
+                        f'Epoch[{int(used_epoch)}] Loss={average_loss:.6f}, P={step_metric_precision:.6f}, R={step_metric_recall:.6f}, F1={step_metric_f1_score:.6f}, Speed = {samples_per_seconds:.1f} samples/s, {progress:.2f} %, ETA = {eta_datetime_string}')
+
+                    # reset for next logging
+                    model.reset_metrics()
+                    loss_sum = 0.0
+                    loss_count = 0
+                    used_sample_sum = 0
+                    last_time = current_time
+
+            offset.assign_add(slice_size)
+            print(f'Saving checkpoint ... ({datetime.datetime.now()})')
+            manager.save()
+
+        used_epoch.assign_add(1)
+        random_seed.assign_add(1)
+        offset.assign(0)
+
+        if int(used_epoch) % export_model_per_epoch == 0:
+            print('Saving model ... (per epoch {export_model_per_epoch})')
+            export_path = os.path.join(
+                project_path, f'model-{model_type}.h5.e{int(used_epoch)}')
+            model.save(export_path, include_optimizer=False, save_format='h5')
+
+    print('Saving model ...')
+    model_path = os.path.join(
+        project_path, f'model-{model_type}.h5')
+
+    # tf.keras.experimental.export_saved_model throw exception now
+    # see https://github.com/tensorflow/tensorflow/issues/27112
+    model.save(model_path, include_optimizer=False)
+
+    print('Training is complete.')
+    print(
+        f'used_epoch={int(used_epoch)}, used_minibatch={int(used_minibatch)}, used_sample={int(used_sample)}')

DeepDanbooru/deepdanbooru/data/__init__.py

@@ -0,0 +1,29 @@
+from typing import Any, Union
+
+import six
+import tensorflow as tf
+
+import deepdanbooru as dd
+
+from .dataset import load_image_records, load_tags
+from .dataset_wrapper import DatasetWrapper
+
+
+def load_image_for_evaluate(
+        input_: Union[str, six.BytesIO], width: int, height: int, normalize: bool = True
+) -> Any:
+    if isinstance(input_, six.BytesIO):
+        image_raw = input_.getvalue()
+    else:
+        image_raw = tf.io.read_file(input_)
+    image = tf.io.decode_png(image_raw, channels=3)
+
+    image = tf.image.resize(
+        image, size=(height, width), method=tf.image.ResizeMethod.AREA, preserve_aspect_ratio=True)
+    image = image.numpy()  # EagerTensor to np.array
+    image = dd.image.transform_and_pad_image(image, width, height)
+
+    if normalize:
+        image = image / 255.0
+
+    return image

DeepDanbooru/deepdanbooru/data/dataset.py

@@ -0,0 +1,40 @@
+import os
+import sqlite3
+
+
+def load_tags(tags_path):
+    with open(tags_path, 'r') as tags_stream:
+        tags = [tag for tag in (tag.strip() for tag in tags_stream) if tag]
+        return tags
+
+
+def load_image_records(sqlite_path, minimum_tag_count):
+    if not os.path.exists(sqlite_path):
+        raise Exception(f'SQLite database is not exists : {sqlite_path}')
+
+    connection = sqlite3.connect(sqlite_path)
+    connection.row_factory = sqlite3.Row
+    cursor = connection.cursor()
+
+    image_folder_path = os.path.join(os.path.dirname(sqlite_path), 'images')
+
+    cursor.execute(
+        "SELECT md5, file_ext, tag_string FROM posts WHERE (file_ext = 'png' OR file_ext = 'jpg' OR file_ext = 'jpeg') AND (tag_count_general >= ?) ORDER BY id",
+        (minimum_tag_count,))
+
+    rows = cursor.fetchall()
+
+    image_records = []
+
+    for row in rows:
+        md5 = row['md5']
+        extension = row['file_ext']
+        image_path = os.path.join(
+            image_folder_path, md5[0:2], f'{md5}.{extension}')
+        tag_string = row['tag_string']
+
+        image_records.append((image_path, tag_string))
+
+    connection.close()
+
+    return image_records

DeepDanbooru/deepdanbooru/data/dataset_wrapper.py

@@ -0,0 +1,98 @@
+import random
+
+import numpy as np
+import tensorflow as tf
+
+import deepdanbooru as dd
+
+
+class DatasetWrapper:
+    """
+    Wrapper class for data pipelining/augmentation.
+    """
+
+    def __init__(self, inputs, tags, width, height, scale_range, rotation_range, shift_range):
+        self.inputs = inputs
+        self.width = width
+        self.height = height
+        self.scale_range = scale_range
+        self.rotation_range = rotation_range
+        self.shift_range = shift_range
+        self.tag_all_array = np.array(tags)
+
+    def get_dataset(self, minibatch_size):
+        dataset = tf.data.Dataset.from_tensor_slices(self.inputs)
+        dataset = dataset.map(
+            self.map_load_image, num_parallel_calls=tf.data.experimental.AUTOTUNE)
+        dataset = dataset.apply(tf.data.experimental.ignore_errors())
+        dataset = dataset.map(
+            self.map_transform_image_and_label, num_parallel_calls=tf.data.experimental.AUTOTUNE)
+        dataset = dataset.batch(minibatch_size)
+        dataset = dataset.prefetch(
+            buffer_size=tf.data.experimental.AUTOTUNE)
+        # dataset = dataset.apply(
+        #    tf.data.experimental.prefetch_to_device('/device:GPU:0'))
+
+        return dataset
+
+    def map_load_image(self, image_path, tag_string):
+        image_raw = tf.io.read_file(image_path)
+        image = tf.io.decode_png(image_raw, channels=3)
+
+        if self.scale_range:
+            pre_scale = self.scale_range[1]
+        else:
+            pre_scale = 1.0
+
+        size = (int(self.height * pre_scale), int(self.width * pre_scale))
+
+        image = tf.image.resize(
+            image, size=size, method=tf.image.ResizeMethod.AREA, preserve_aspect_ratio=True)
+
+        return (image, tag_string)
+
+    def map_transform_image_and_label(self, image, tag_string):
+        return tf.py_function(self.map_transform_image_and_label_py, (image, tag_string), (tf.float32, tf.float32))
+
+    def map_transform_image_and_label_py(self, image, tag_string):
+        # transform image
+        image = image.numpy()
+
+        if self.scale_range:
+            scale = random.uniform(
+                self.scale_range[0], self.scale_range[1]) * (1.0 / self.scale_range[1])
+        else:
+            scale = None
+
+        if self.rotation_range:
+            rotation = random.uniform(
+                self.rotation_range[0], self.rotation_range[1])
+        else:
+            rotation = None
+
+        if self.shift_range:
+            shift_x = random.uniform(self.shift_range[0], self.shift_range[1])
+            shift_y = random.uniform(self.shift_range[0], self.shift_range[1])
+            shift = (shift_x, shift_y)
+        else:
+            shift = None
+
+        image = dd.image.transform_and_pad_image(
+            image=image,
+            target_width=self.width,
+            target_height=self.height,
+            rotation=rotation,
+            scale=scale,
+            shift=shift)
+
+        image = image / 255.0  # normalize to 0~1
+        # image = image.astype(np.float32)
+
+        # transform tag
+        tag_string = tag_string.numpy().decode()
+        tag_array = np.array(tag_string.split(' '))
+
+        labels = np.where(np.isin(self.tag_all_array,
+                                  tag_array), 1, 0).astype(np.float32)
+
+        return (image, labels)

DeepDanbooru/deepdanbooru/extra/__init__.py

@@ -0,0 +1,18 @@
+import re
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    """
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    """
+    return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def natural_sorted(iterable):
+    return sorted(iterable, key=natural_keys)

DeepDanbooru/deepdanbooru/gradcam.py

@@ -0,0 +1,48 @@
+import tensorflow as tf
+import numpy as np
+
+# gpus = tf.config.experimental.list_physical_devices('GPU')
+
+# if gpus:
+#     try:
+#         for gpu in gpus:
+#             tf.config.experimental.set_memory_growth(gpu, True)
+#     except RuntimeError as e:
+#         print(e)
+
+
+def grad(y, x):
+    V = tf.keras.layers.Lambda(lambda z: tf.gradients(
+        z[0], z[1]), output_shape=[1])([y, x])
+    return V
+
+
+def grad_cam_test(model, x, some_variable):
+    fixed_input = model.inputs
+    fixed_output = tf.keras.layers.Lambda(lambda z: tf.keras.backend.gradients(
+        z[0], z[1]), output_shape=[2])([model.inputs[0], model.outputs[0]])
+
+    grad_model = tf.keras.Model(inputs=fixed_input, outputs=fixed_output)
+
+    return grad_model.predict(x)
+
+
+def run_test():
+    # Generate sample model
+    x = tf.keras.Input(shape=(2))
+    y = tf.keras.layers.Dense(2)(x)
+    model = tf.keras.Model(inputs=x, outputs=y)
+    target = np.array([[1.0, 2.0]], dtype=np.float32)
+
+    # Calculate gradient using numpy array
+    input_numpy = np.array([[0.0, 0.0]])
+    grad_output_numpy = grad_cam_test(model, input_numpy, target)
+    print(f'numpy: {grad_output_numpy}')
+
+    # Calculate gradient using tf.Variable
+    input_variable = tf.constant([[0.0, 0.0]])
+    grad_output_variable = grad_cam_test(model, input_variable, target)
+    print(f'variable: {grad_output_variable}')
+
+
+run_test()

DeepDanbooru/deepdanbooru/image/__init__.py

@@ -0,0 +1,56 @@
+import math
+
+import numpy as np
+import skimage.transform
+
+
+def calculate_image_scale(source_width, source_height, target_width, target_height):
+    """
+    Calculate scale for image resizing while preserving aspect ratio.
+    """
+    if source_width == target_width and source_height == target_height:
+        return 1.0
+
+    source_ratio = source_width / source_height
+    target_ratio = target_width / target_height
+
+    if target_ratio < source_ratio:
+        scale = target_width / source_width
+    else:
+        scale = target_height / source_height
+
+    return scale
+
+
+def transform_and_pad_image(image, target_width, target_height, scale=None, rotation=None, shift=None, order=1, mode='edge'):
+    """
+    Transform image and pad by edge pixles.
+    """
+    image_width = image.shape[1]
+    image_height = image.shape[0]
+    image_array = image
+
+    # centerize
+    t = skimage.transform.AffineTransform(
+        translation=(-image_width * 0.5, -image_height * 0.5))
+
+    if scale:
+        t += skimage.transform.AffineTransform(scale=(scale, scale))
+
+    if rotation:
+        radian = (rotation / 180.0) * math.pi
+        t += skimage.transform.AffineTransform(rotation=radian)
+
+    t += skimage.transform.AffineTransform(
+        translation=(target_width * 0.5, target_height * 0.5))
+
+    if shift:
+        t += skimage.transform.AffineTransform(
+            translation=(target_width * shift[0], target_height * shift[1]))
+
+    warp_shape = (target_height, target_width)
+
+    image_array = skimage.transform.warp(
+        image_array, (t).inverse, output_shape=warp_shape, order=order, mode=mode)
+
+    return image_array

DeepDanbooru/deepdanbooru/io/__init__.py

@@ -0,0 +1,28 @@
+import json
+import os
+from pathlib import Path
+
+
+def serialize_as_json(target_object, path, encoding='utf-8'):
+    with open(path, 'w', encoding=encoding) as stream:
+        stream.write(json.dumps(target_object, indent=4, ensure_ascii=False))
+
+
+def deserialize_from_json(path, encoding='utf-8'):
+    with open(path, 'r', encoding=encoding) as stream:
+        return json.loads(stream.read())
+
+
+def try_create_directory(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+
+def get_file_paths_in_directory(path, patterns):
+    return [str(file_path) for pattern in patterns for file_path in Path(path).rglob(pattern)]
+
+
+def get_image_file_paths_recursive(folder_path, patterns_string):
+    patterns = patterns_string.split(',')
+
+    return get_file_paths_in_directory(folder_path, patterns)

DeepDanbooru/deepdanbooru/model/__init__.py

@@ -0,0 +1,8 @@
+import deepdanbooru.model.layers
+import deepdanbooru.model.losses
+
+from .resnet import create_resnet_152
+from .resnet import create_resnet_custom_v1
+from .resnet import create_resnet_custom_v2
+from .resnet import create_resnet_custom_v3
+from .resnet import create_resnet_custom_v4

DeepDanbooru/deepdanbooru/model/layers/__init__.py

@@ -0,0 +1,60 @@
+import tensorflow as tf
+
+
+def conv(x, filters, kernel_size, strides=(1, 1), padding='same', initializer='he_normal'):
+    c = tf.keras.layers.Conv2D(
+        filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, kernel_initializer=initializer, use_bias=False)(x)
+
+    return c
+
+
+def conv_bn(x, filters, kernel_size, strides=(1, 1), padding='same', initializer='he_normal', bn_gamma_initializer='ones'):
+    c = conv(x, filters=filters, kernel_size=kernel_size,
+             strides=strides, padding=padding, initializer=initializer)
+
+    c_bn = tf.keras.layers.BatchNormalization(
+        gamma_initializer=bn_gamma_initializer)(c)
+
+    return c_bn
+
+
+def conv_bn_relu(x, filters, kernel_size, strides=(1, 1), padding='same', initializer='he_normal', bn_gamma_initializer='ones'):
+    c_bn = conv_bn(x, filters=filters, kernel_size=kernel_size, strides=strides, padding=padding,
+                   initializer=initializer, bn_gamma_initializer=bn_gamma_initializer)
+
+    return tf.keras.layers.Activation('relu')(c_bn)
+
+
+def conv_gap(x, output_filters, kernel_size=(1, 1)):
+    x = conv(x, filters=output_filters, kernel_size=kernel_size)
+    x = tf.keras.layers.GlobalAveragePooling2D()(x)
+
+    return x
+
+
+def repeat_blocks(x, block_delegate, count, **kwargs):
+    assert count >= 0
+
+    for _ in range(count):
+        x = block_delegate(x, **kwargs)
+    return x
+
+
+def squeeze_excitation(x, reduction=16):
+    """
+    Squeeze-Excitation layer from https://arxiv.org/abs/1709.01507
+    """
+    output_filters = x.shape[-1]
+
+    assert output_filters // reduction > 0
+
+    s = x
+
+    s = tf.keras.layers.GlobalAveragePooling2D()(s)
+    s = tf.keras.layers.Dense(
+        output_filters//reduction, activation='relu')(x)
+    s = tf.keras.layers.Dense(
+        output_filters, activation='sigmoid')(x)
+    x = tf.keras.layers.Multiply()([x, s])
+
+    return x

DeepDanbooru/deepdanbooru/model/losses/__init__.py

@@ -0,0 +1,24 @@
+import tensorflow as tf
+
+
+def focal_loss(alpha=0.25, gamma=2.0, epsilon=1e-6):
+    def loss(y_true, y_pred):
+        value = -alpha * y_true * tf.math.pow(1.0 - y_pred, gamma) * tf.math.log(y_pred + epsilon) - (
+            1.0 - alpha) * (1.0 - y_true) * tf.math.pow(y_pred, gamma) * tf.math.log(1.0 - y_pred + epsilon)
+
+        return tf.math.reduce_sum(value)
+
+    return loss
+
+
+def binary_crossentropy(epsilon=1e-6):
+    def loss(y_true, y_pred):
+        clipped_y_pred = tf.clip_by_value(y_pred, epsilon, tf.float32.max)
+        clipped_y_pred_nega = tf.clip_by_value(
+            1.0 - y_pred, epsilon, tf.float32.max)
+        value = (-y_true) * tf.math.log(clipped_y_pred) - \
+            (1.0 - y_true) * tf.math.log(clipped_y_pred_nega)
+
+        return tf.math.reduce_sum(value)
+
+    return loss

DeepDanbooru/deepdanbooru/model/resnet.py

@@ -0,0 +1,190 @@
+import numpy as np
+import tensorflow as tf
+import deepdanbooru as dd
+
+
+def resnet_bottleneck_block(x, output_filters, inter_filters, activation=True, se=False):
+    c1 = dd.model.layers.conv_bn_relu(x, inter_filters, (1, 1))
+    c2 = dd.model.layers.conv_bn_relu(c1, inter_filters, (3, 3))
+    c3 = dd.model.layers.conv_bn(
+        c2, output_filters, (1, 1), bn_gamma_initializer='zeros')
+
+    if se:
+        c3 = dd.model.layers.squeeze_excitation(c3)
+
+    p = tf.keras.layers.Add()([c3, x])
+
+    if activation:
+        return tf.keras.layers.Activation('relu')(p)
+    else:
+        return p
+
+
+def resnet_bottleneck_inc_block(x, output_filters, inter_filters, strides1x1=(1, 1), strides2x2=(2, 2), se=False):
+    c1 = dd.model.layers.conv_bn_relu(
+        x, inter_filters, (1, 1), strides=strides1x1)
+    c2 = dd.model.layers.conv_bn_relu(
+        c1, inter_filters, (3, 3), strides=strides2x2)
+    c3 = dd.model.layers.conv_bn(
+        c2, output_filters, (1, 1), bn_gamma_initializer='zeros')
+
+    if se:
+        c3 = dd.model.layers.squeeze_excitation(c3)
+
+    strides = np.multiply(strides1x1, strides2x2)
+    s = dd.model.layers.conv_bn(
+        x, output_filters, (1, 1), strides=strides)  # shortcut
+
+    p = tf.keras.layers.Add()([c3, s])
+
+    return tf.keras.layers.Activation('relu')(p)
+
+
+def resnet_original_bottleneck_model(x, filter_sizes, repeat_sizes, final_pool=True, se=False):
+    """
+    https://github.com/Microsoft/CNTK/blob/master/Examples/Image/Classification/ResNet/Python/resnet_models.py
+    """
+    assert len(filter_sizes) == len(repeat_sizes)
+
+    x = dd.model.layers.conv_bn_relu(
+        x, filter_sizes[0] // 4, (7, 7), strides=(2, 2))
+    x = tf.keras.layers.MaxPool2D((3, 3), strides=(
+        2, 2), padding='same')(x)
+
+    for i in range(len(repeat_sizes)):
+        x = resnet_bottleneck_inc_block(
+            x=x,
+            output_filters=filter_sizes[i],
+            inter_filters=filter_sizes[i] // 4,
+            strides2x2=(2, 2) if i > 0 else (1, 1),
+            se=se)
+        x = dd.model.layers.repeat_blocks(
+            x=x,
+            block_delegate=resnet_bottleneck_block,
+            count=repeat_sizes[i],
+            output_filters=filter_sizes[i],
+            inter_filters=filter_sizes[i] // 4,
+            se=se)
+
+    if final_pool:
+        x = tf.keras.layers.AveragePooling2D((7, 7), name='ap_final')(x)
+
+    return x
+
+
+def resnet_longterm_bottleneck_model(x, filter_sizes, repeat_sizes, final_pool=True, se=False):
+    """
+    Add long-term shortcut.
+    """
+    assert len(filter_sizes) == len(repeat_sizes)
+
+    x = dd.model.layers.conv_bn_relu(
+        x, filter_sizes[0] // 4, (7, 7), strides=(2, 2))
+    x = tf.keras.layers.MaxPool2D((3, 3), strides=(
+        2, 2), padding='same')(x)
+
+    for i in range(len(repeat_sizes)):
+        x = resnet_bottleneck_inc_block(
+            x=x,
+            output_filters=filter_sizes[i],
+            inter_filters=filter_sizes[i] // 4,
+            strides2x2=(2, 2) if i > 0 else (1, 1),
+            se=se)
+        x_1 = dd.model.layers.repeat_blocks(
+            x=x,
+            block_delegate=resnet_bottleneck_block,
+            count=repeat_sizes[i] - 1,
+            output_filters=filter_sizes[i],
+            inter_filters=filter_sizes[i] // 4,
+            se=se)
+        x_1 = resnet_bottleneck_block(
+            x_1,
+            output_filters=filter_sizes[i],
+            inter_filters=filter_sizes[i] // 4,
+            activation=False,
+            se=se)
+
+        x = tf.keras.layers.Add()([x_1, x])  # long-term shortcut
+        x = tf.keras.layers.Activation('relu')(x)
+
+    if final_pool:
+        x = tf.keras.layers.AveragePooling2D((7, 7), name='ap_final')(x)
+
+    return x
+
+
+def create_resnet_152(x, output_dim):
+    """
+    Original ResNet-152 Model.
+    """
+    filter_sizes = [256, 512, 1024, 2048]
+    repeat_sizes = [2, 7, 35, 2]
+
+    x = resnet_original_bottleneck_model(
+        x, filter_sizes=filter_sizes, repeat_sizes=repeat_sizes)
+
+    x = tf.keras.layers.Flatten()(x)
+    x = tf.keras.layers.Dense(output_dim)(x)
+    x = tf.keras.layers.Activation('sigmoid')(x)
+
+    return x
+
+
+def create_resnet_custom_v1(x, output_dim):
+    """
+    DeepDanbooru web (until 2019/04/20)
+    Short, wide
+    """
+    filter_sizes = [256, 512, 1024, 2048, 4096]
+    repeat_sizes = [2, 7, 35, 2, 2]
+
+    x = resnet_original_bottleneck_model(
+        x, filter_sizes=filter_sizes, repeat_sizes=repeat_sizes, final_pool=False)
+
+    x = dd.model.layers.conv_gap(x, output_dim)
+    x = tf.keras.layers.Activation('sigmoid')(x)
+
+    return x
+
+
+def create_resnet_custom_v2(x, output_dim):
+    """
+    Experimental (blazing-deep network)
+    Deep, narrow
+    """
+    filter_sizes = [256, 512, 1024, 1024, 1024, 2048]
+    repeat_sizes = [2, 7, 40, 16, 16, 6]
+
+    x = resnet_original_bottleneck_model(
+        x, filter_sizes=filter_sizes, repeat_sizes=repeat_sizes, final_pool=False)
+
+    x = dd.model.layers.conv_gap(x, output_dim)
+    x = tf.keras.layers.Activation('sigmoid')(x)
+
+    return x
+
+
+def create_resnet_custom_v3(x, output_dim):
+    filter_sizes = [256, 512, 1024, 1024, 2048, 4096]
+    repeat_sizes = [2, 7, 19, 19, 2, 2]
+
+    x = resnet_original_bottleneck_model(
+        x, filter_sizes=filter_sizes, repeat_sizes=repeat_sizes, final_pool=False)
+
+    x = dd.model.layers.conv_gap(x, output_dim)
+    x = tf.keras.layers.Activation('sigmoid')(x)
+
+    return x
+
+
+def create_resnet_custom_v4(x, output_dim):
+    filter_sizes = [256, 512, 1024, 1024, 1024, 2048]
+    repeat_sizes = [2, 7, 10, 10, 10, 2]
+
+    x = resnet_original_bottleneck_model(
+        x, filter_sizes=filter_sizes, repeat_sizes=repeat_sizes, final_pool=False)
+
+    x = dd.model.layers.conv_gap(x, output_dim)
+    x = tf.keras.layers.Activation('sigmoid')(x)
+
+    return x

DeepDanbooru/deepdanbooru/project/__init__.py

@@ -0,0 +1,4 @@
+from .project import DEFAULT_PROJECT_CONTEXT
+from .project import load_project
+from .project import load_model_from_project
+from .project import load_tags_from_project

DeepDanbooru/deepdanbooru/project/project.py

@@ -0,0 +1,50 @@
+import os
+import deepdanbooru as dd
+import tensorflow as tf
+
+DEFAULT_PROJECT_CONTEXT = {
+    'image_width': 299,
+    'image_height': 299,
+    'database_path': None,
+    'minimum_tag_count': 20,
+    'model': 'resnet_custom_v2',
+    'minibatch_size': 32,
+    'epoch_count': 10,
+    'export_model_per_epoch': 10,
+    'checkpoint_frequency_mb': 200,
+    'console_logging_frequency_mb': 10,
+    'optimizer': 'adam',
+    'learning_rate': 0.001,
+    'rotation_range': [0.0, 360.0],
+    'scale_range': [0.9, 1.1],
+    'shift_range': [-0.1, 0.1]
+}
+
+
+def load_project(project_path):
+    project_context_path = os.path.join(project_path, 'project.json')
+    project_context = dd.io.deserialize_from_json(project_context_path)
+    tags = dd.data.load_tags_from_project(project_path)
+
+    model_type = project_context['model']
+    model_path = os.path.join(project_path, f'model-{model_type}.h5')
+    model = tf.keras.models.load_model(model_path)
+
+    return project_context, model, tags
+
+
+def load_model_from_project(project_path, compile_model=True):
+    project_context_path = os.path.join(project_path, 'project.json')
+    project_context = dd.io.deserialize_from_json(project_context_path)
+
+    model_type = project_context['model']
+    model_path = os.path.join(project_path, f'model-{model_type}.h5')
+    model = tf.keras.models.load_model(model_path, compile=compile_model)
+
+    return model
+
+
+def load_tags_from_project(project_path):
+    tags_path = os.path.join(project_path, 'tags.txt')
+
+    return dd.data.load_tags(tags_path)

DeepDanbooru/repeat-run.bat

@@ -0,0 +1,15 @@
+@echo off
+:start
+set RETRY_COUNT=0
+set TF_CPP_MIN_LOG_LEVEL=2
+
+:run
+call %* || goto :wait
+pause
+exit /B 0
+
+:wait
+echo %date% %time% Retry %RETRY_COUNT%
+set /a RETRY_COUNT=%RETRY_COUNT%+1
+timeout /t 3 /nobreak
+goto :run

DeepDanbooru/setup.cfg

@@ -0,0 +1,14 @@
+[mypy]
+ignore_missing_imports = True
+[flake8]
+max-line-length = 260
+exclude =
+  .tox,
+  venv
+ignore =
+  # E226 missing whitespace around arithmetic operator
+  # F401 'X' imported but unused
+  # W503 line break before binary operator
+  E226,
+  F401,
+  W503,

DeepDanbooru/setup.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+import re
+
+import setuptools
+
+with open("README.md", "r", encoding='utf-8') as fh:
+    long_description = fh.read()
+
+
+with open('deepdanbooru/__main__.py', encoding='utf-8') as f:
+    version = re.search('__version__ = \'([^\']+)\'', f.read()).group(1)  # type: ignore
+
+
+install_requires = [
+]
+tensorflow_pkg = 'tensorflow>=2.1.0'
+
+setuptools.setup(
+    name="deepdanbooru",
+    version=version,
+    author="Kichang Kim",
+    author_email="admin@kanotype.net",
+    description="DeepDanbooru is AI based multi-label girl image classification system, "
+    "implemented by using TensorFlow.",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/KichangKim/DeepDanbooru",
+    packages=setuptools.find_packages(),
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: MIT License",
+        "Operating System :: OS Independent",
+    ],
+    python_requires='>=3.6',
+    install_requires=install_requires,
+    extras_require={
+        'tensorflow': [tensorflow_pkg],
+        'test': ['pytest', 'flake8', 'mypy']
+    },
+    entry_points={
+        "console_scripts": [
+            "deepdanbooru=deepdanbooru.__main__:main",
+        ]
+    },
+)

Dockerfile

@@ -9,6 +9,8 @@ RUN pip config set global.index-url https://pypi.mirrors.ustc.edu.cn/simple
 RUN pip install -r requirements.txt
 
 COPY . /app
+RUN cd DeepDanbooru && python setup.py install
+
 ENV HF_TOKEN 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
 
 CMD ["python", "app.py"]

requirements.txt

@@ -1,3 +1,9 @@
 pillow>=9.0.0
 tensorflow>=2.7.0
-git+https://github.com/KichangKim/DeepDanbooru@v3-20200915-sgd-e30#egg=deepdanbooru
+Click>=7.0
+numpy>=1.16.2
+scikit-image>=0.15.0
+tensorflow>=2.1.0
+requests>=2.22.0
+six>=1.13.0
+#git+https://github.com/KichangKim/DeepDanbooru@v3-20200915-sgd-e30#egg=deepdanbooru