71 lines
2.7 KiB
Python
71 lines
2.7 KiB
Python
import os
|
|
import logging
|
|
import traceback
|
|
import tensorflow as tf
|
|
from tensorflow.keras import layers, models
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
|
|
|
def build_clean_model(num_classes, img_size=(224, 224)):
|
|
base_model = tf.keras.applications.EfficientNetV2B0(
|
|
input_shape=img_size + (3,),
|
|
include_top=False,
|
|
weights=None,
|
|
)
|
|
inputs = tf.keras.Input(shape=img_size + (3,))
|
|
x = base_model(inputs, training=False)
|
|
x = layers.Conv2D(512, (3, 3), padding='same', activation='swish')(x)
|
|
x = layers.BatchNormalization()(x)
|
|
x = layers.MaxPooling2D((2, 2))(x)
|
|
x = layers.Dropout(0.2)(x)
|
|
x = layers.Conv2D(256, (3, 3), padding='same', activation='swish')(x)
|
|
x = layers.BatchNormalization()(x)
|
|
x = layers.GlobalAveragePooling2D()(x)
|
|
x = layers.Dropout(0.3)(x)
|
|
x = layers.Dense(1024, activation='swish')(x)
|
|
x = layers.BatchNormalization()(x)
|
|
x = layers.Dropout(0.4)(x)
|
|
outputs = layers.Dense(num_classes, activation='softmax', dtype='float32')(x)
|
|
return models.Model(inputs, outputs)
|
|
|
|
logging.info("=== EXPORT STARTED ===")
|
|
try:
|
|
MODEL_KERAS_PATH = "best_model/best_model.keras"
|
|
OUTPUT_DIR = "model"
|
|
saved_model_dir = os.path.join(OUTPUT_DIR, "saved_model")
|
|
tflite_path = os.path.join(OUTPUT_DIR, "model.tflite")
|
|
|
|
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
|
|
|
if not os.path.exists(MODEL_KERAS_PATH):
|
|
raise FileNotFoundError(f"Model file not found at {MODEL_KERAS_PATH}")
|
|
|
|
logging.info(f"Loading trained weights from {MODEL_KERAS_PATH}...")
|
|
original_model = tf.keras.models.load_model(MODEL_KERAS_PATH, compile=False)
|
|
|
|
logging.info("Building clean architecture...")
|
|
clean_model = build_clean_model(num_classes=original_model.output_shape[-1])
|
|
clean_model.set_weights(original_model.get_weights())
|
|
logging.info("Weights cloned successfully.")
|
|
|
|
logging.info(f"Exporting to SavedModel format at: {saved_model_dir}...")
|
|
tf.saved_model.save(clean_model, saved_model_dir)
|
|
logging.info("SavedModel export completed successfully.")
|
|
|
|
logging.info(f"Converting to TFLite format at: {tflite_path}...")
|
|
converter = tf.lite.TFLiteConverter.from_keras_model(clean_model)
|
|
converter.target_spec.supported_ops = [
|
|
tf.lite.OpsSet.TFLITE_BUILTINS,
|
|
tf.lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
converter.optimizations = [tf.lite.Optimize.DEFAULT]
|
|
tflite_model = converter.convert()
|
|
with open(tflite_path, "wb") as f:
|
|
f.write(tflite_model)
|
|
logging.info("TFLite conversion completed successfully.")
|
|
logging.info("=== EXPORT COMPLETED ===")
|
|
|
|
except Exception:
|
|
logging.error("EXPORT FAILED")
|
|
logging.error(traceback.format_exc())
|