Python | OpenCV TensorFlow/Keras 微調範例

# clean_dataset.py
import os
from PIL import Image

SUPPORTED_FORMATS = {"JPEG", "PNG", "GIF", "BMP"}
DATA_ROOT = "data/cat_dog"

removed = 0
for split in ["train", "val"]:
    split_dir = f"{DATA_ROOT}/{split}"
    if not os.path.isdir(split_dir):
        continue
    for cls in os.listdir(split_dir):
        folder = os.path.join(split_dir, cls)
        if not os.path.isdir(folder):
            continue
        for fname in os.listdir(folder):
            fpath = os.path.join(folder, fname)
            try:
                with Image.open(fpath) as img:
                    fmt = img.format          # 讀副檔名前先取實際格式
                    if fmt not in SUPPORTED_FORMATS:
                        raise ValueError(f"不支援的格式：{fmt}")
                    img.load()               # 實際解碼，抓損壞的檔案
            except Exception as e:
                os.remove(fpath)
                print(f"移除：{fpath}（{e}）")
                removed += 1

print(f"清理完成，共移除 {removed} 個檔案")

# train.py
import os
import json
import tensorflow as tf
from tensorflow.keras import layers, Model
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping, ReduceLROnPlateau

# ── 超參數設定 ──────────────────────────────
IMG_SIZE    = 224
BATCH_SIZE  = 32
NUM_CLASSES = 2                            # 依實際類別數修改
TASK_NAME   = "cat_dog"                    # 換任務只改這一行
DATA_DIR    = f"data/{TASK_NAME}"
SAVE_DIR    = f"models/mobilenetv2/{TASK_NAME}"
# ────────────────────────────────────────────

os.makedirs(SAVE_DIR, exist_ok=True)
SAVE_PATH = f"{SAVE_DIR}/best_model.keras"

# ── 資料載入 ────────────────────────────────
train_ds = tf.keras.utils.image_dataset_from_directory(
    f"{DATA_DIR}/train",
    image_size=(IMG_SIZE, IMG_SIZE),
    batch_size=BATCH_SIZE,
    shuffle=True,
    seed=42
)

val_ds = tf.keras.utils.image_dataset_from_directory(
    f"{DATA_DIR}/val",
    image_size=(IMG_SIZE, IMG_SIZE),
    batch_size=BATCH_SIZE,
    shuffle=False,
    seed=42
)

class_names = train_ds.class_names
print(f"類別：{class_names}")

# 儲存類別設定
config = {
    "model": "mobilenetv2",
    "num_classes": NUM_CLASSES,
    "classes": class_names
}
with open(f"{SAVE_DIR}/config.json", "w") as f:
    json.dump(config, f, ensure_ascii=False, indent=2)

# 效能優化：prefetch
AUTOTUNE = tf.data.AUTOTUNE
train_ds = train_ds.prefetch(buffer_size=AUTOTUNE)
val_ds   = val_ds.prefetch(buffer_size=AUTOTUNE)

# ── 資料增強層 ──────────────────────────────
data_augmentation = tf.keras.Sequential([
    layers.RandomFlip("horizontal"),
    layers.RandomRotation(0.1),
    layers.RandomZoom(0.1),
    layers.RandomBrightness(0.2),
], name="data_augmentation")

# ── 模型建立 ────────────────────────────────
base_model = MobileNetV2(
    input_shape=(IMG_SIZE, IMG_SIZE, 3),
    include_top=False,   # 移除原本的 1000 類分類頭，只保留特徵提取器
    weights="imagenet"
)
base_model.trainable = False  # 第一階段：凍結所有預訓練層

inputs  = tf.keras.Input(shape=(IMG_SIZE, IMG_SIZE, 3))
x       = data_augmentation(inputs)
x       = tf.keras.applications.mobilenet_v2.preprocess_input(x)
x       = base_model(x, training=False)
x       = layers.GlobalAveragePooling2D()(x)
x       = layers.Dropout(0.3)(x)
outputs = layers.Dense(NUM_CLASSES, activation="softmax")(x)
model   = Model(inputs, outputs)

model.summary()

# ── 第一階段：只訓練分類頭 ──────────────────
model.compile(
    optimizer=tf.keras.optimizers.Adam(1e-3),
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"]
)

callbacks_phase1 = [
    ModelCheckpoint(SAVE_PATH, monitor="val_accuracy",
                    save_best_only=True, verbose=1),
    EarlyStopping(monitor="val_accuracy", patience=5,
                  restore_best_weights=True, verbose=1),
]

print("\n── 第一階段訓練（Feature Extraction）──")
history1 = model.fit(
    train_ds, validation_data=val_ds,
    epochs=10, callbacks=callbacks_phase1
)

# ── 第二階段：解凍後幾層進行 Fine-Tuning ───
base_model.trainable = True

# 只解凍最後 30 層
fine_tune_at = len(base_model.layers) - 30
for layer in base_model.layers[:fine_tune_at]:
    layer.trainable = False

# 使用極小學習率
model.compile(
    optimizer=tf.keras.optimizers.Adam(1e-5),
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"]
)

callbacks_phase2 = [
    ModelCheckpoint(SAVE_PATH, monitor="val_accuracy",
                    save_best_only=True, verbose=1),
    EarlyStopping(monitor="val_accuracy", patience=5,
                  restore_best_weights=True, verbose=1),
    ReduceLROnPlateau(monitor="val_loss", factor=0.5,
                      patience=3, verbose=1),
]

print("\n── 第二階段訓練（Fine-Tuning）──")
history2 = model.fit(
    train_ds, validation_data=val_ds,
    epochs=10, callbacks=callbacks_phase2
)

print(f"\n訓練完成，模型已儲存至 {SAVE_PATH}")

# inference.py
import json
import tensorflow as tf
import numpy as np
import cv2

MODEL_DIR  = "models/mobilenetv2/cat_dog"         # 與 train.py 的 SAVE_DIR 對應
IMAGE_PATH = "data/cat_dog/val/cat/000001.jpg"    # 替換為實際測試圖片路徑
IMG_SIZE   = 224

# 從 config.json 讀取類別資訊，不需要手動填寫
with open(f"{MODEL_DIR}/config.json") as f:
    config = json.load(f)
CLASS_NAMES = config["classes"]

model = tf.keras.models.load_model(f"{MODEL_DIR}/best_model.keras")

img = cv2.imread(IMAGE_PATH)
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_resized = cv2.resize(img_rgb, (IMG_SIZE, IMG_SIZE))

x = np.expand_dims(img_resized, axis=0).astype(np.float32)
preds = model.predict(x, verbose=0)
pred_class = np.argmax(preds[0])
conf = preds[0][pred_class]

print(f"預測類別：{CLASS_NAMES[pred_class]}，信心度：{conf:.4f}")