Merge from upstream

.gitignore

@@ -5,6 +5,8 @@ __pycache__/
 test.py
 /logs
 /uploads
+/build
+/deepdanbooru.egg-info
 *.jpg
 *.jpeg
 *.png

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,16 +57,18 @@ 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/
 │   ├── 00/
 │   │   ├── 00000000000000000000000000000000.jpg
 │   │   ├── ...
 │   ├── 01/
+│   │   ├── 01000000000000000000000000000000.jpg
 │   │   ├── ...
 │   └── ff/
+│       ├── ff000000000000000000000000000000.jpg
 │       ├── ...
 └── my-dataset.sqlite
 ```

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

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,24 +172,34 @@ 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]",
     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,
+    target_paths, # I guess its this one
     project_path,
     model_path,
     tags_path,
     threshold,
     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()

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
+

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}")

deepdanbooru/commands/download_tags.py

@@ -42,6 +42,7 @@ def download_category_tags(
             request_url,
             params=parameters,
         )
+
         response_json = response.json()
 
         response_tags = [

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,14 +41,15 @@ def evaluate_image(
 
 
 def evaluate(
-    target_paths,
+    target_paths, #this
     project_path,
     model_path,
     tags_path,
     threshold,
     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()

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,25 +292,25 @@ 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
     model.save(model_path, include_optimizer=False)
 
     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.")

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'))
 

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

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