175 lines
6.8 KiB
Python
175 lines
6.8 KiB
Python
import os
|
|
import json
|
|
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")
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
def cbam_block(x, ratio=8, name="cbam"):
|
|
"""Convolutional Block Attention Module — lightweight foreground attention."""
|
|
channels = x.shape[-1]
|
|
|
|
# Channel attention
|
|
avg_pool = layers.GlobalAveragePooling2D()(x)
|
|
max_pool = layers.GlobalMaxPooling2D()(x)
|
|
ca = layers.Dense(channels // ratio, activation="swish", name=f"{name}_ca1")(avg_pool)
|
|
ca = layers.Dense(channels, activation="sigmoid", name=f"{name}_ca2")(ca)
|
|
ca2 = layers.Dense(channels // ratio, activation="swish", name=f"{name}_ca3")(max_pool)
|
|
ca2 = layers.Dense(channels, activation="sigmoid", name=f"{name}_ca4")(ca2)
|
|
ca_out = layers.Add(name=f"{name}_ca_add")([ca, ca2])
|
|
ca_out = layers.Reshape((1, 1, channels), name=f"{name}_ca_reshape")(ca_out)
|
|
x = layers.Multiply(name=f"{name}_ca_mul")([x, ca_out])
|
|
|
|
# Spatial attention
|
|
from keras import ops
|
|
avg_sp = ops.mean(x, axis=-1, keepdims=True)
|
|
max_sp = ops.max(x, axis=-1, keepdims=True)
|
|
sp = layers.Concatenate(name=f"{name}_sa_cat")([avg_sp, max_sp])
|
|
sp = layers.Conv2D(1, 7, padding="same", activation="sigmoid", name=f"{name}_sa_conv")(sp)
|
|
x = layers.Multiply(name=f"{name}_sa_mul")([x, sp])
|
|
return x
|
|
|
|
|
|
def build_clean_model(num_classes, target_size=(224, 224)):
|
|
"""Build the production architecture: CBAM + lightweight head, outputting raw logits.
|
|
|
|
Mirrors the notebook's build_model() exactly so set_weights() maps correctly.
|
|
"""
|
|
base_model = tf.keras.applications.EfficientNetV2B0(
|
|
input_shape=target_size + (3,),
|
|
include_top=False,
|
|
weights=None,
|
|
)
|
|
base_model.trainable = False
|
|
|
|
inputs = tf.keras.Input(shape=(None, None, 3), name="input")
|
|
x = layers.Resizing(target_size[0], target_size[1], interpolation="bilinear",
|
|
name="resize_input")(inputs)
|
|
x = layers.GaussianNoise(0.05, name="gauss_noise")(x)
|
|
x = base_model(x, training=False)
|
|
# CBAM attention — focus on leaf regions, ignore background
|
|
x = cbam_block(x, ratio=8, name="cbam")
|
|
x = layers.GlobalAveragePooling2D(name="gap")(x)
|
|
x = layers.Dropout(0.3, name="drop_gap")(x)
|
|
x = layers.Dense(512, activation="swish", name="dense_head")(x)
|
|
x = layers.BatchNormalization(name="bn_head")(x)
|
|
x = layers.Dropout(0.4, name="drop_head")(x)
|
|
# Raw logits (no softmax) — temperature scaling applied at inference
|
|
outputs = layers.Dense(num_classes, activation="linear", dtype="float32", name="logits")(x)
|
|
return models.Model(inputs, outputs)
|
|
|
|
|
|
def _read_labels_for_classes():
|
|
"""Detect number of classes from model/labels.json or best_model/calibration.json."""
|
|
for path in ["model/labels.json", "best_model/calibration.json"]:
|
|
if os.path.exists(path):
|
|
with open(path) as f:
|
|
meta = json.load(f)
|
|
labels = meta.get("labels")
|
|
if labels:
|
|
return len(labels)
|
|
log.warning("Could not detect num_classes from metadata; defaulting to 4.")
|
|
return 4
|
|
|
|
|
|
log.info("=== EXPORT STARTED (v3.0) ===")
|
|
try:
|
|
CKPT_DIR = "best_model"
|
|
WEIGHTS_PATH = os.path.join(CKPT_DIR, "model.weights.h5")
|
|
MODEL_KERAS_PATH = os.path.join(CKPT_DIR, "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)
|
|
|
|
# Use weights H5 as primary source (portable, no Lambda serialization issues)
|
|
if not os.path.exists(WEIGHTS_PATH):
|
|
raise FileNotFoundError(
|
|
f"Weights file not found at {WEIGHTS_PATH}. "
|
|
"Run the notebook Cell 31 (Save Model) first to generate it."
|
|
)
|
|
|
|
log.info(f"Detecting model configuration...")
|
|
num_classes = _read_labels_for_classes()
|
|
|
|
log.info(f"Building export architecture ({num_classes} classes)...")
|
|
clean_model = build_clean_model(num_classes=num_classes)
|
|
|
|
log.info(f"Loading trained weights from {WEIGHTS_PATH}...")
|
|
clean_model.load_weights(WEIGHTS_PATH)
|
|
log.info("Weights loaded successfully.")
|
|
|
|
|
|
# ─── Export SavedModel (raw logits) ───
|
|
log.info(f"Exporting SavedModel (raw logits) at: {saved_model_dir}...")
|
|
tf.saved_model.save(clean_model, saved_model_dir)
|
|
log.info("SavedModel export completed successfully.")
|
|
|
|
# ─── Export TFLite (INT8 quantization) ───
|
|
log.info(f"Converting to TFLite INT8 at: {tflite_path}...")
|
|
|
|
# Representative dataset for INT8 quantization
|
|
def representative_dataset():
|
|
val_dir = "dataset_split/val"
|
|
if not os.path.exists(val_dir):
|
|
log.warning("Validation dir not found; skipping representative dataset.")
|
|
return
|
|
ds = tf.keras.utils.image_dataset_from_directory(
|
|
val_dir, shuffle=True, batch_size=1, image_size=(224, 224)
|
|
)
|
|
for images, _ in ds.take(200):
|
|
yield [tf.cast(images, tf.float32)]
|
|
|
|
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]
|
|
if os.path.exists("dataset_split/val"):
|
|
converter.representative_dataset = representative_dataset
|
|
|
|
tflite_model = converter.convert()
|
|
with open(tflite_path, "wb") as f:
|
|
f.write(tflite_model)
|
|
log.info("TFLite conversion completed successfully.")
|
|
|
|
# ─── Export model metadata ───
|
|
# Load labels and calibration from training output
|
|
labels_path = os.path.join(OUTPUT_DIR, "labels.json")
|
|
cal_path = os.path.join("best_model", "calibration.json")
|
|
|
|
labels_meta = {"labels": None, "temperature": 1.0, "conf_threshold_high": 0.70,
|
|
"conf_threshold_low": 0.45}
|
|
|
|
if os.path.exists(labels_path):
|
|
with open(labels_path) as f:
|
|
labels_meta.update(json.load(f))
|
|
|
|
if os.path.exists(cal_path):
|
|
with open(cal_path) as f:
|
|
cal = json.load(f)
|
|
labels_meta["temperature"] = cal.get("temperature", 1.0)
|
|
|
|
labels_meta["version"] = "3.0"
|
|
labels_meta["architecture"] = "EfficientNetV2B0 + CBAM + Dense(512)"
|
|
labels_meta["output_type"] = "logits"
|
|
labels_meta["input_range"] = [0, 255]
|
|
labels_meta["input_size"] = [224, 224]
|
|
labels_meta["preprocessing"] = "resize_bilinear_224x224_no_normalization"
|
|
|
|
with open(labels_path, "w") as f:
|
|
json.dump(labels_meta, f, indent=2)
|
|
log.info(f"Labels + calibration metadata saved to {labels_path}")
|
|
|
|
log.info("=== EXPORT COMPLETED (v3.0) ===")
|
|
|
|
except Exception:
|
|
log.error("EXPORT FAILED")
|
|
log.error(traceback.format_exc())
|