diff --git a/.gitignore b/.gitignore index e845071..9c5607b 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,8 @@ __pycache__/ test.py /logs /uploads +/build +/deepdanbooru.egg-info *.jpg *.jpeg *.png diff --git a/README.md b/README.md index 9d72d06..37e0fb7 100644 --- a/README.md +++ b/README.md @@ -1,19 +1,19 @@ # DeepDanbooru -[![Python](https://img.shields.io/badge/python-3.6-green)](https://www.python.org/doc/versions/) +[![Python](https://img.shields.io/badge/python-3.11-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-20200915-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://dev.kanotype.net:8003/deepdanbooru/). ## Requirements -DeepDanbooru is written by Python 3.7. Following packages are need to be installed. -- tensorflow>=2.7.0 -- tensorflow-io>=2.22.0 -- Click>=7.0 -- numpy>=1.16.2 -- requests>=2.22.0 -- scikit-image>=0.15.0 -- six>=1.13.0 +DeepDanbooru is written by Python 3.11. Following packages are need to be installed. +- Click>=8.1.7 +- numpy>=1.26.4 +- requests>=2.32.3 +- scikit-image>=0.24.0 +- six>=1.16.0 +- tensorflow>=2.17.0 +- tensorflow-io>=0.31.0 Or just use `requirements.txt`. ``` @@ -38,9 +38,9 @@ To install it with tensorflow, add `tensorflow` extra package. ``` > 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. +3. Prepare tag list. If you want to use latest tags, use following command. It downloads tag from Danbooru server. (Need Danbooru account and API key) ``` -> deepdanbooru download-tags [your_project_folder] +> deepdanbooru download-tags [your_project_folder] --username [your_danbooru_account] --api-key [your_danbooru_api_key] ``` 4. (Option) Filtering dataset. If you want to train with optional tags (rating and score), you should convert it as system tags. ``` @@ -57,7 +57,7 @@ To install it with tensorflow, add `tensorflow` extra package. ``` ## 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. +DeepDanbooru uses following folder structure for input dataset. SQLite file can be any name, but must be located in same folder to `images` folder. All of image files are located in sub-folder which named first 2 characters of its filename. ``` MyDataset/ ├── images/ @@ -65,8 +65,10 @@ MyDataset/ │ │ ├── 00000000000000000000000000000000.jpg │ │ ├── ... │ ├── 01/ +│ │ ├── 01000000000000000000000000000000.jpg │ │ ├── ... │ └── ff/ +│ ├── ff000000000000000000000000000000.jpg │ ├── ... └── my-dataset.sqlite ``` diff --git a/deepdanbooru/__init__.py b/deepdanbooru/__init__.py index 6cb06ea..674fc93 100644 --- a/deepdanbooru/__init__.py +++ b/deepdanbooru/__init__.py @@ -1,3 +1,9 @@ +import warnings +warnings.filterwarnings("ignore") + +import os +os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' + import deepdanbooru.commands import deepdanbooru.data import deepdanbooru.extra diff --git a/deepdanbooru/__main__.py b/deepdanbooru/__main__.py index 68b3388..35c4a2f 100644 --- a/deepdanbooru/__main__.py +++ b/deepdanbooru/__main__.py @@ -2,8 +2,10 @@ from typing_extensions import Required import click - +import warnings +import os import deepdanbooru as dd +import tensorflow.lite as tflite __version__ = "1.0.0" @@ -170,6 +172,12 @@ def grad_cam(project_path, target_path, output_path, threshold): 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( + "--save-txt", + default=False, + is_flag=True, + help="Enable this option to save tags to a txt file with the same filename.", +) @click.option( "--folder-filters", default="*.[Pp][Nn][Gg],*.[Jj][Pp][Gg],*.[Jj][Pp][Ee][Gg],*.[Gg][Ii][Ff]", @@ -177,7 +185,7 @@ def grad_cam(project_path, target_path, output_path, threshold): ) @click.option("--verbose", default=False, is_flag=True) def evaluate( - target_paths, + target_paths, # I guess its this one project_path, model_path, tags_path, @@ -185,9 +193,13 @@ def evaluate( allow_gpu, compile_model, allow_folder, + save_txt, folder_filters, verbose, ): + if verbose: + warnings.filterwarnings("always") + os.environ['TF_CPP_MIN_LOG_LEVEL'] = '0' dd.commands.evaluate( target_paths, project_path, @@ -197,10 +209,39 @@ def evaluate( allow_gpu, compile_model, allow_folder, + save_txt, folder_filters, verbose, ) +@main.command("conv2tflite", help="Convert saved model into tflite model.") +@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( + "--save-path", + type=click.Path(resolve_path=True, file_okay=True, dir_okay=False), +) +@click.option("--optimize-default", default=True, is_flag=True) +@click.option("--optimize-experimental-sparsity", default=False, is_flag=True) +@click.option("--verbose", default=False, is_flag=True) +def conv2tflite(project_path, model_path, save_path, optimize_default, optimize_experimental_sparsity, verbose): + if verbose: + warnings.filterwarnings("always") + os.environ['TF_CPP_MIN_LOG_LEVEL'] = '0' + if not optimize_default and not optimize_experimental_sparsity: + raise Exception("optimization method must be specified") + op = [] + if optimize_default: op = [tflite.Optimize.DEFAULT] + if optimize_experimental_sparsity: op.append(tflite.Optimize.EXPERIMENTAL_SPARSITY) + dd.commands.convert_to_tflite_from_from_saved_model(project_path, model_path, save_path, op, verbose=verbose) + if __name__ == "__main__": main() diff --git a/deepdanbooru/commands/__init__.py b/deepdanbooru/commands/__init__.py index 1c4f572..ceab469 100644 --- a/deepdanbooru/commands/__init__.py +++ b/deepdanbooru/commands/__init__.py @@ -1,7 +1,9 @@ +from .convert_to_tflite import convert_to_tflite_from_from_saved_model 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 +from .grad_cam import grad_cam +from .make_training_database import make_training_database +from .train_project import train_project + diff --git a/deepdanbooru/commands/convert_to_tflite.py b/deepdanbooru/commands/convert_to_tflite.py new file mode 100644 index 0000000..59d34f1 --- /dev/null +++ b/deepdanbooru/commands/convert_to_tflite.py @@ -0,0 +1,40 @@ +from typing import List +import tensorflow as tf +import deepdanbooru as dd + +# optimizations: see in tf.lite.Optimize +def convert_to_tflite_from_from_saved_model( + project_path: str, model_path: str, save_path: str, + optimizations: List[tf.lite.Optimize] = [tf.lite.Optimize.DEFAULT], + verbose: bool = False +): + if not model_path and not project_path: + raise Exception("You must provide project path or model path.") + + if not save_path: + raise Exception("You must provide a path to save tflite model.") + + if model_path: + if verbose: + print(f"Loading model from {model_path} ...") + model = tf.keras.models.load_model(model_path) + else: + if verbose: + print(f"Loading model from project {project_path} ...") + model = dd.project.load_model_from_project(project_path) + + if verbose: + print("Converting ...") + + converter = tf.lite.TFLiteConverter.from_keras_model(model) + converter.optimizations = optimizations + tflite_model = converter.convert() + + if verbose: + print("Saving ...") + + with open(save_path, "wb") as f: + f.write(tflite_model) + + if verbose: + print(f"Converted model has been saved to {save_path}") diff --git a/deepdanbooru/commands/download_tags.py b/deepdanbooru/commands/download_tags.py index 4b4e053..cdcedef 100644 --- a/deepdanbooru/commands/download_tags.py +++ b/deepdanbooru/commands/download_tags.py @@ -42,6 +42,7 @@ def download_category_tags( request_url, params=parameters, ) + response_json = response.json() response_tags = [ diff --git a/deepdanbooru/commands/evaluate.py b/deepdanbooru/commands/evaluate.py index a6b6bd8..94936e8 100644 --- a/deepdanbooru/commands/evaluate.py +++ b/deepdanbooru/commands/evaluate.py @@ -6,6 +6,17 @@ import deepdanbooru as dd +def save_txt_file(txt_path, list): + last_index = len(list)-1 + last_tag = list[last_index] + with open(txt_path, 'w') as writer: + for i in list: + if last_tag == i: + writer.write(i) + writer.close() + else: + writer.write(i + ", ") + print("Saved text file.") def evaluate_image( image_input: Union[str, six.BytesIO], model: Any, tags: List[str], threshold: float @@ -30,7 +41,7 @@ def evaluate_image( def evaluate( - target_paths, + target_paths, #this project_path, model_path, tags_path, @@ -38,6 +49,7 @@ def evaluate( allow_gpu, compile_model, allow_folder, + save_txt, folder_filters, verbose, ): @@ -83,8 +95,12 @@ def evaluate( tags = dd.project.load_tags_from_project(project_path) for image_path in target_image_paths: - print(f"Tags of {image_path}:") + print(f"Tags of {image_path}:") #yup! + if save_txt: tag_list = [] for tag, score in evaluate_image(image_path, model, tags, threshold): print(f"({score:05.3f}) {tag}") - + if save_txt: tag_list.append(tag) + if save_txt: + txt_file_path = str(os.path.splitext(image_path)[0]) + ".txt" + save_txt_file(txt_file_path, tag_list) print() diff --git a/deepdanbooru/commands/train_project.py b/deepdanbooru/commands/train_project.py index c7260b6..a8c6e2c 100644 --- a/deepdanbooru/commands/train_project.py +++ b/deepdanbooru/commands/train_project.py @@ -206,7 +206,7 @@ def train_project(project_path, source_model): 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} ...") + tf.print(f"Learning rate is changed to", optimizer.learning_rate, "...") while int(offset) < epoch_size: image_records_slice = image_records[ @@ -231,7 +231,7 @@ def train_project(project_path, source_model): sample_count = x_train.shape[0] step_result = model.train_on_batch( - x_train, y_train, reset_metrics=False + x_train, y_train ) used_minibatch.assign_add(1) @@ -292,17 +292,17 @@ def train_project(project_path, source_model): if export_model_per_epoch == 0 or int(used_epoch) % export_model_per_epoch == 0: print(f"Saving model ... (per epoch {export_model_per_epoch})") export_path = os.path.join( - project_path, f"model-{model_type}.h5.e{int(used_epoch)}" + project_path, f"model-{model_type}.e{int(used_epoch)}.keras" ) - model.save(export_path, include_optimizer=False, save_format="h5") + model.save(export_path, include_optimizer=False) if use_mixed_precision: export_model_as_float32( - model_float32, checkpoint_path, export_path + ".float32.h5" + model_float32, checkpoint_path, export_path + ".float32.keras" ) print("Saving model ...") - model_path = os.path.join(project_path, f"model-{model_type}.h5") + model_path = os.path.join(project_path, f"model-{model_type}.keras") # tf.keras.experimental.export_saved_model throw exception now # see https://github.com/tensorflow/tensorflow/issues/27112 @@ -310,7 +310,7 @@ def train_project(project_path, source_model): if use_mixed_precision: export_model_as_float32( - model_float32, checkpoint_path, model_path + ".float32.h5" + model_float32, checkpoint_path, model_path + ".float32.keras" ) print("Training is complete.") diff --git a/deepdanbooru/data/dataset_wrapper.py b/deepdanbooru/data/dataset_wrapper.py index 2bcf27c..94e9ab0 100644 --- a/deepdanbooru/data/dataset_wrapper.py +++ b/deepdanbooru/data/dataset_wrapper.py @@ -26,15 +26,15 @@ def __init__( 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 + self.map_load_image, num_parallel_calls=tf.data.AUTOTUNE ) - dataset = dataset.apply(tf.data.experimental.ignore_errors()) + dataset = dataset.ignore_errors() dataset = dataset.map( self.map_transform_image_and_label, - num_parallel_calls=tf.data.experimental.AUTOTUNE, + num_parallel_calls=tf.data.AUTOTUNE, ) dataset = dataset.batch(minibatch_size) - dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE) + dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE) # dataset = dataset.apply( # tf.data.experimental.prefetch_to_device('/device:GPU:0')) diff --git a/deepdanbooru/project/project.py b/deepdanbooru/project/project.py index b56a8d8..271e3cb 100644 --- a/deepdanbooru/project/project.py +++ b/deepdanbooru/project/project.py @@ -29,7 +29,11 @@ def load_project(project_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_path = os.path.join(project_path, f"model-{model_type}.keras") + + if not os.path.isfile(model_path): + 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 @@ -40,7 +44,11 @@ def load_model_from_project(project_path, compile_model=True): 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_path = os.path.join(project_path, f"model-{model_type}.keras") + + if not os.path.isfile(model_path): + 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 diff --git a/requirements.txt b/requirements.txt index d9078da..38fff21 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,7 @@ -Click>=7.0 -numpy>=1.16.2 -scikit-image>=0.15.0 -tensorflow>=2.7.0 -tensorflow-io>=0.22.0 -requests>=2.22.0 -six>=1.13.0 +Click>=8.1.7 +numpy>=1.26.4 +scikit-image>=0.24.0 +tensorflow>=2.17.0 +tensorflow-io>=0.31.0 +requests>=2.32.3 +six>=1.16.0