From 38889296447214bec1db70fab3335d2019f9cf88 Mon Sep 17 00:00:00 2001 From: jianjiang Date: Wed, 26 Apr 2023 18:25:25 +0800 Subject: [PATCH] reslove --- DeepDanbooru/.gitignore | 15 ++ DeepDanbooru/.travis.yml | 12 + DeepDanbooru/LICENSE | 21 ++ DeepDanbooru/README.md | 99 ++++++++ DeepDanbooru/deepdanbooru/__init__.py | 8 + DeepDanbooru/deepdanbooru/__main__.py | 89 +++++++ .../deepdanbooru/commands/__init__.py | 7 + .../deepdanbooru/commands/create_project.py | 15 ++ .../deepdanbooru/commands/download_tags.py | 162 +++++++++++++ .../deepdanbooru/commands/evaluate.py | 77 ++++++ .../deepdanbooru/commands/evaluate_project.py | 51 ++++ .../deepdanbooru/commands/grad_cam.py | 111 +++++++++ .../commands/make_training_database.py | 122 ++++++++++ .../deepdanbooru/commands/train_project.py | 222 ++++++++++++++++++ DeepDanbooru/deepdanbooru/data/__init__.py | 29 +++ DeepDanbooru/deepdanbooru/data/dataset.py | 40 ++++ .../deepdanbooru/data/dataset_wrapper.py | 98 ++++++++ DeepDanbooru/deepdanbooru/extra/__init__.py | 18 ++ DeepDanbooru/deepdanbooru/gradcam.py | 48 ++++ DeepDanbooru/deepdanbooru/image/__init__.py | 56 +++++ DeepDanbooru/deepdanbooru/io/__init__.py | 28 +++ DeepDanbooru/deepdanbooru/model/__init__.py | 8 + .../deepdanbooru/model/layers/__init__.py | 60 +++++ .../deepdanbooru/model/losses/__init__.py | 24 ++ DeepDanbooru/deepdanbooru/model/resnet.py | 190 +++++++++++++++ DeepDanbooru/deepdanbooru/project/__init__.py | 4 + DeepDanbooru/deepdanbooru/project/project.py | 50 ++++ DeepDanbooru/deepdanbooru/train/__init__.py | 0 DeepDanbooru/repeat-run.bat | 15 ++ DeepDanbooru/setup.cfg | 14 ++ DeepDanbooru/setup.py | 46 ++++ Dockerfile | 2 + requirements.txt | 8 +- 33 files changed, 1748 insertions(+), 1 deletion(-) create mode 100644 DeepDanbooru/.gitignore create mode 100644 DeepDanbooru/.travis.yml create mode 100644 DeepDanbooru/LICENSE create mode 100644 DeepDanbooru/README.md create mode 100644 DeepDanbooru/deepdanbooru/__init__.py create mode 100644 DeepDanbooru/deepdanbooru/__main__.py create mode 100644 DeepDanbooru/deepdanbooru/commands/__init__.py create mode 100644 DeepDanbooru/deepdanbooru/commands/create_project.py create mode 100644 DeepDanbooru/deepdanbooru/commands/download_tags.py create mode 100644 DeepDanbooru/deepdanbooru/commands/evaluate.py create mode 100644 DeepDanbooru/deepdanbooru/commands/evaluate_project.py create mode 100644 DeepDanbooru/deepdanbooru/commands/grad_cam.py create mode 100644 DeepDanbooru/deepdanbooru/commands/make_training_database.py create mode 100644 DeepDanbooru/deepdanbooru/commands/train_project.py create mode 100644 DeepDanbooru/deepdanbooru/data/__init__.py create mode 100644 DeepDanbooru/deepdanbooru/data/dataset.py create mode 100644 DeepDanbooru/deepdanbooru/data/dataset_wrapper.py create mode 100644 DeepDanbooru/deepdanbooru/extra/__init__.py create mode 100644 DeepDanbooru/deepdanbooru/gradcam.py create mode 100644 DeepDanbooru/deepdanbooru/image/__init__.py create mode 100644 DeepDanbooru/deepdanbooru/io/__init__.py create mode 100644 DeepDanbooru/deepdanbooru/model/__init__.py create mode 100644 DeepDanbooru/deepdanbooru/model/layers/__init__.py create mode 100644 DeepDanbooru/deepdanbooru/model/losses/__init__.py create mode 100644 DeepDanbooru/deepdanbooru/model/resnet.py create mode 100644 DeepDanbooru/deepdanbooru/project/__init__.py create mode 100644 DeepDanbooru/deepdanbooru/project/project.py create mode 100644 DeepDanbooru/deepdanbooru/train/__init__.py create mode 100644 DeepDanbooru/repeat-run.bat create mode 100644 DeepDanbooru/setup.cfg create mode 100644 DeepDanbooru/setup.py diff --git a/DeepDanbooru/.gitignore b/DeepDanbooru/.gitignore new file mode 100644 index 0000000..e845071 --- /dev/null +++ b/DeepDanbooru/.gitignore @@ -0,0 +1,15 @@ +.vs/ +.vscode/ +__pycache__/ +/test* +test.py +/logs +/uploads +*.jpg +*.jpeg +*.png +*.bmp +*.gif +*.h5 +*.txt +*.json diff --git a/DeepDanbooru/.travis.yml b/DeepDanbooru/.travis.yml new file mode 100644 index 0000000..ee4c93d --- /dev/null +++ b/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 . diff --git a/DeepDanbooru/LICENSE b/DeepDanbooru/LICENSE new file mode 100644 index 0000000..d977d82 --- /dev/null +++ b/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. diff --git a/DeepDanbooru/README.md b/DeepDanbooru/README.md new file mode 100644 index 0000000..868d6bc --- /dev/null +++ b/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 +... +``` diff --git a/DeepDanbooru/deepdanbooru/__init__.py b/DeepDanbooru/deepdanbooru/__init__.py new file mode 100644 index 0000000..6cb06ea --- /dev/null +++ b/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 diff --git a/DeepDanbooru/deepdanbooru/__main__.py b/DeepDanbooru/deepdanbooru/__main__.py new file mode 100644 index 0000000..3778f76 --- /dev/null +++ b/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() diff --git a/DeepDanbooru/deepdanbooru/commands/__init__.py b/DeepDanbooru/deepdanbooru/commands/__init__.py new file mode 100644 index 0000000..1c4f572 --- /dev/null +++ b/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 diff --git a/DeepDanbooru/deepdanbooru/commands/create_project.py b/DeepDanbooru/deepdanbooru/commands/create_project.py new file mode 100644 index 0000000..feeadfd --- /dev/null +++ b/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})') diff --git a/DeepDanbooru/deepdanbooru/commands/download_tags.py b/DeepDanbooru/deepdanbooru/commands/download_tags.py new file mode 100644 index 0000000..a48377f --- /dev/null +++ b/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.') diff --git a/DeepDanbooru/deepdanbooru/commands/evaluate.py b/DeepDanbooru/deepdanbooru/commands/evaluate.py new file mode 100644 index 0000000..8071f76 --- /dev/null +++ b/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() diff --git a/DeepDanbooru/deepdanbooru/commands/evaluate_project.py b/DeepDanbooru/deepdanbooru/commands/evaluate_project.py new file mode 100644 index 0000000..1926137 --- /dev/null +++ b/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() diff --git a/DeepDanbooru/deepdanbooru/commands/grad_cam.py b/DeepDanbooru/deepdanbooru/commands/grad_cam.py new file mode 100644 index 0000000..cb4eca1 --- /dev/null +++ b/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('/', '_'))) diff --git a/DeepDanbooru/deepdanbooru/commands/make_training_database.py b/DeepDanbooru/deepdanbooru/commands/make_training_database.py new file mode 100644 index 0000000..1ad5eb0 --- /dev/null +++ b/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() diff --git a/DeepDanbooru/deepdanbooru/commands/train_project.py b/DeepDanbooru/deepdanbooru/commands/train_project.py new file mode 100644 index 0000000..9814045 --- /dev/null +++ b/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)}') diff --git a/DeepDanbooru/deepdanbooru/data/__init__.py b/DeepDanbooru/deepdanbooru/data/__init__.py new file mode 100644 index 0000000..c63418f --- /dev/null +++ b/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 diff --git a/DeepDanbooru/deepdanbooru/data/dataset.py b/DeepDanbooru/deepdanbooru/data/dataset.py new file mode 100644 index 0000000..32b5d50 --- /dev/null +++ b/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 diff --git a/DeepDanbooru/deepdanbooru/data/dataset_wrapper.py b/DeepDanbooru/deepdanbooru/data/dataset_wrapper.py new file mode 100644 index 0000000..327e80b --- /dev/null +++ b/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) diff --git a/DeepDanbooru/deepdanbooru/extra/__init__.py b/DeepDanbooru/deepdanbooru/extra/__init__.py new file mode 100644 index 0000000..5270ee2 --- /dev/null +++ b/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) diff --git a/DeepDanbooru/deepdanbooru/gradcam.py b/DeepDanbooru/deepdanbooru/gradcam.py new file mode 100644 index 0000000..0b59a50 --- /dev/null +++ b/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() diff --git a/DeepDanbooru/deepdanbooru/image/__init__.py b/DeepDanbooru/deepdanbooru/image/__init__.py new file mode 100644 index 0000000..c3f2cc4 --- /dev/null +++ b/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 diff --git a/DeepDanbooru/deepdanbooru/io/__init__.py b/DeepDanbooru/deepdanbooru/io/__init__.py new file mode 100644 index 0000000..824d2b2 --- /dev/null +++ b/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) diff --git a/DeepDanbooru/deepdanbooru/model/__init__.py b/DeepDanbooru/deepdanbooru/model/__init__.py new file mode 100644 index 0000000..fb66416 --- /dev/null +++ b/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 diff --git a/DeepDanbooru/deepdanbooru/model/layers/__init__.py b/DeepDanbooru/deepdanbooru/model/layers/__init__.py new file mode 100644 index 0000000..729b68f --- /dev/null +++ b/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 diff --git a/DeepDanbooru/deepdanbooru/model/losses/__init__.py b/DeepDanbooru/deepdanbooru/model/losses/__init__.py new file mode 100644 index 0000000..7e5468c --- /dev/null +++ b/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 diff --git a/DeepDanbooru/deepdanbooru/model/resnet.py b/DeepDanbooru/deepdanbooru/model/resnet.py new file mode 100644 index 0000000..4db1bce --- /dev/null +++ b/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 diff --git a/DeepDanbooru/deepdanbooru/project/__init__.py b/DeepDanbooru/deepdanbooru/project/__init__.py new file mode 100644 index 0000000..92566e1 --- /dev/null +++ b/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 diff --git a/DeepDanbooru/deepdanbooru/project/project.py b/DeepDanbooru/deepdanbooru/project/project.py new file mode 100644 index 0000000..0a8092b --- /dev/null +++ b/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) diff --git a/DeepDanbooru/deepdanbooru/train/__init__.py b/DeepDanbooru/deepdanbooru/train/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/DeepDanbooru/repeat-run.bat b/DeepDanbooru/repeat-run.bat new file mode 100644 index 0000000..7266427 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/DeepDanbooru/setup.cfg b/DeepDanbooru/setup.cfg new file mode 100644 index 0000000..0890fe0 --- /dev/null +++ b/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, diff --git a/DeepDanbooru/setup.py b/DeepDanbooru/setup.py new file mode 100644 index 0000000..2ac92ee --- /dev/null +++ b/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", + ] + }, +) diff --git a/Dockerfile b/Dockerfile index 5f49fcb..f6ba6c7 100644 --- a/Dockerfile +++ b/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"] diff --git a/requirements.txt b/requirements.txt index 595418e..630bfec 100644 --- a/requirements.txt +++ b/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