mirror of
https://github.com/KichangKim/DeepDanbooru.git
synced 2026-10-27 12:29:09 +08:00
61 lines
1.8 KiB
Python
61 lines
1.8 KiB
Python
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,
|
|
"loss": "binary_crossentropy",
|
|
"optimizer": "adam",
|
|
"learning_rate": 0.001,
|
|
"rotation_range": [0.0, 360.0],
|
|
"scale_range": [0.9, 1.1],
|
|
"shift_range": [-0.1, 0.1],
|
|
"mixed_precision": False,
|
|
}
|
|
|
|
|
|
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}.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
|
|
|
|
|
|
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}.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
|
|
|
|
|
|
def load_tags_from_project(project_path):
|
|
tags_path = os.path.join(project_path, "tags.txt")
|
|
|
|
return dd.data.load_tags(tags_path)
|