reslove
Build-Deploy-Actions
Details
Build-Deploy-Actions
Details
This commit is contained in:
parent
3eddca3e36
commit
3888929644
|
|
@ -0,0 +1,15 @@
|
|||
.vs/
|
||||
.vscode/
|
||||
__pycache__/
|
||||
/test*
|
||||
test.py
|
||||
/logs
|
||||
/uploads
|
||||
*.jpg
|
||||
*.jpeg
|
||||
*.png
|
||||
*.bmp
|
||||
*.gif
|
||||
*.h5
|
||||
*.txt
|
||||
*.json
|
||||
|
|
@ -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 .
|
||||
|
|
@ -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.
|
||||
|
|
@ -0,0 +1,99 @@
|
|||
# DeepDanbooru
|
||||
[](https://www.python.org/doc/versions/)
|
||||
[](https://opensource.org/licenses/MIT)
|
||||
[](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
|
||||
...
|
||||
```
|
||||
|
|
@ -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
|
||||
|
|
@ -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()
|
||||
|
|
@ -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
|
||||
|
|
@ -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})')
|
||||
|
|
@ -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.')
|
||||
|
|
@ -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()
|
||||
|
|
@ -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()
|
||||
|
|
@ -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('/', '_')))
|
||||
|
|
@ -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()
|
||||
|
|
@ -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)}')
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
@ -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)
|
||||
|
|
@ -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()
|
||||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
@ -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,
|
||||
|
|
@ -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",
|
||||
]
|
||||
},
|
||||
)
|
||||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue