app.py

"""
FastAPI Inference Server — Latency Experiment
Paper: Analisis Komparatif Pretrained CNN untuk Klasifikasi Penyakit Daun Tomat
Author: Mohammad Wisam Wiraghina

Struktur folder yang diperlukan di HF Spaces:
    app.py                          ← file ini
    requirements.txt
    models/
        mobilenetv3small.h5
        mobilenetv3large.h5
        efficientnetb0.h5
        efficientnetv2b0.h5
        resnet50.h5
        inceptionv3.h5
        xception.h5
"""

import os
import time
import logging
from contextlib import asynccontextmanager
from typing import Dict

import numpy as np
from PIL import Image
import io

import tensorflow as tf
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse

# ── Logging ────────────────────────────────────────────────────────────────────
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_DIR = os.path.join(ROOT_DIR, "models")

# ── Konstanta ──────────────────────────────────────────────────────────────────
CLASS_NAMES = [
    "Bacterial_Spot",
    "Early_Blight",
    "Late_Blight",
    "Leaf_Mold",
    "Septoria_Leaf_Spot",
    "Spider_Mites",
    "Target_Spot",
    "Yellow_Leaf_Curl_Virus",
    "Mosaic_Virus",
    "Healthy",
]

IMAGE_SIZE = (224, 224)

# Mapping nama model → file .h5 dan fungsi preprocess_input
MODEL_CONFIG: Dict[str, dict] = {
    "mobilenetv3small": {
        "path": os.path.join(MODEL_DIR, "MobileNetV3Small_best.h5"),
        "preprocess": "mobilenetv3",
    },
    "mobilenetv3large": {
        "path": os.path.join(MODEL_DIR, "MobileNetV3Large_best.h5"),
        "preprocess": "mobilenetv3",
    },
    "efficientnetb0": {
        "path": os.path.join(MODEL_DIR, "EfficientNetB0_best.h5"),
        "preprocess": "efficientnet",
    },
    "efficientnetv2b0": {
        "path": os.path.join(MODEL_DIR, "EfficientNetV2B0_best.h5"),
        "preprocess": "efficientnetv2",
    },
    "resnet50": {
        "path": os.path.join(MODEL_DIR, "ResNet50_best.h5"),
        "preprocess": "resnet50",
    },
    "inceptionv3": {
        "path": os.path.join(MODEL_DIR, "InceptionV3_best.h5"),
        "preprocess": "inception",   # normalize ke [-1, 1]
    },
    "xception": {
        "path": os.path.join(MODEL_DIR, "Xception_best.h5"),
        "preprocess": "inception",   # normalize ke [-1, 1]
    },
}

# ── Preprocess dispatcher ──────────────────────────────────────────────────────
def preprocess_image(img_array: np.ndarray, mode: str) -> np.ndarray:
    """
    Menerapkan preprocess_input sesuai backbone.
    img_array shape: (1, 224, 224, 3), dtype float32, range [0, 255]
    """
    if mode == "mobilenetv3":
        return tf.keras.applications.mobilenet_v3.preprocess_input(img_array)
    elif mode == "efficientnet":
        return tf.keras.applications.efficientnet.preprocess_input(img_array)
    elif mode == "efficientnetv2":
        return tf.keras.applications.efficientnet_v2.preprocess_input(img_array)
    elif mode == "resnet50":
        return tf.keras.applications.resnet50.preprocess_input(img_array)
    elif mode == "inception":
        # InceptionV3 & Xception → [-1, 1]
        return tf.keras.applications.inception_v3.preprocess_input(img_array)
    else:
        raise ValueError(f"Unknown preprocess mode: {mode}")

# ── Compat layer: handle model version mismatch ──────────────────────────────
# Model disimpan dengan Keras versi baru (Kaggle) yang punya field baru
# seperti quantization_config di Dense. Kita strip field-field yang tidak
# dikenali agar load_model tidak error.
import keras
import h5py

keras.config.enable_unsafe_deserialization()


_OriginalDense = keras.layers.Dense
class _CompatDense(_OriginalDense):
    def __init__(self, *args, quantization_config=None, **kwargs):
        super().__init__(*args, **kwargs)

# Patch juga layer lain yang mungkin punya quantization_config
_OriginalConv2D = keras.layers.Conv2D
class _CompatConv2D(_OriginalConv2D):
    def __init__(self, *args, quantization_config=None, **kwargs):
        super().__init__(*args, **kwargs)

_OriginalBatchNorm = keras.layers.BatchNormalization
class _CompatBatchNorm(_OriginalBatchNorm):
    def __init__(self, *args, quantization_config=None, **kwargs):
        super().__init__(*args, **kwargs)

_OriginalDepthwiseConv2D = keras.layers.DepthwiseConv2D
class _CompatDepthwiseConv2D(_OriginalDepthwiseConv2D):
    def __init__(self, *args, quantization_config=None, **kwargs):
        super().__init__(*args, **kwargs)

CUSTOM_OBJECTS = {
    "Dense": _CompatDense,
    "Conv2D": _CompatConv2D,
    "BatchNormalization": _CompatBatchNorm,
    "DepthwiseConv2D": _CompatDepthwiseConv2D,
}

# ── MobileNetV3 manual loader (bypass Keras 3.x hard_swish bug) ──────────────
MOBILENETV3_MODELS = {"mobilenetv3small", "mobilenetv3large"}

def _load_mobilenetv3_manual(model_name: str, h5_path: str) -> keras.Model:
    """
    Rebuild MobileNetV3 architecture from scratch, then inject weights
    from h5 via h5py. This bypasses the hard_swish deserialization bug
    in Keras 3.x that prevents load_model from working on MobileNetV3.
    """
    from keras.applications import MobileNetV3Small, MobileNetV3Large
    from keras.layers import GlobalAveragePooling2D, Dense, Dropout
    from keras.models import Model

    BackboneFn = MobileNetV3Small if model_name == "mobilenetv3small" else MobileNetV3Large
    backbone_key = "MobileNetV3Small" if model_name == "mobilenetv3small" else "MobileNetV3Large"

    base = BackboneFn(weights="imagenet", include_top=False, input_shape=(224, 224, 3))
    base.trainable = False
    x = base.output
    x = GlobalAveragePooling2D()(x)
    x = Dense(256, activation="relu", name="dense")(x)
    x = Dropout(0.4, name="dropout")(x)
    x = Dense(128, activation="relu", name="dense_1")(x)
    x = Dropout(0.3, name="dropout_1")(x)
    out = Dense(10, activation="softmax", name="predictions")(x)
    model = Model(inputs=base.input, outputs=out)

    # Load weights manually from h5
    f = h5py.File(h5_path, "r")
    mw = f["model_weights"]

    # Classifier head: find dense layer names in h5 (may be dense/dense_1 or dense_2/dense_3)
    h5_dense_names = sorted([n for n in mw.keys()
                             if n.startswith("dense") and "dropout" not in n
                             and "prediction" not in n])
    head_mapping = [
        (model.get_layer("dense"), h5_dense_names[0]),
        (model.get_layer("dense_1"), h5_dense_names[1]),
        (model.get_layer("predictions"), "predictions"),
    ]
    for layer, h5_name in head_mapping:
        g = mw[h5_name][h5_name]
        layer.set_weights([g["kernel"][()], g["bias"][()]])

    # Backbone weights
    backbone_group = mw[backbone_key]
    for layer in base.layers:
        if layer.name in backbone_group:
            lg = backbone_group[layer.name]
            if layer.name in lg:
                lg = lg[layer.name]
            w_list = [lg[wn][()] for wn in sorted(lg.keys())]
            if w_list:
                try:
                    layer.set_weights(w_list)
                except Exception:
                    pass

    f.close()
    logger.info(f"  [manual loader] {model_name}: architecture rebuilt, weights injected")
    return model

# ── Global model registry ──────────────────────────────────────────────────────
# Model di-load sekali saat startup, disimpan di dict ini
LOADED_MODELS: Dict[str, keras.Model] = {}
LOAD_ERRORS: Dict[str, str] = {}  # simpan error untuk diagnostik

# ── Lifespan: load semua model saat startup ────────────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("=== Loading all models ===")
    for model_name, config in MODEL_CONFIG.items():
        model_path = config["path"]
        if not os.path.exists(model_path):
            msg = f"file tidak ditemukan di {model_path}"
            logger.warning(f"[SKIP] {model_name}: {msg}")
            LOAD_ERRORS[model_name] = msg
            continue
        try:
            logger.info(f"Loading {model_name} dari {model_path} ...")
            if model_name in MOBILENETV3_MODELS:
                model = _load_mobilenetv3_manual(model_name, model_path)
            else:
                model = keras.models.load_model(
                    model_path, compile=False, custom_objects=CUSTOM_OBJECTS
                )
            # Warmup internal: satu forward pass dummy untuk inisialisasi CPU graph
            dummy = np.zeros((1, 224, 224, 3), dtype=np.float32)
            model.predict(dummy, verbose=0)
            LOADED_MODELS[model_name] = model
            logger.info(f"  ✓ {model_name} siap ({model.count_params():,} params)")
        except Exception as e:
            import traceback
            err_msg = f"{type(e).__name__}: {e}"
            logger.error(f"  ✗ Gagal load {model_name}: {err_msg}")
            logger.error(traceback.format_exc())
            LOAD_ERRORS[model_name] = err_msg

    logger.info(f"=== {len(LOADED_MODELS)}/{len(MODEL_CONFIG)} model berhasil dimuat ===")
    yield  # server berjalan di sini
    LOADED_MODELS.clear()
    logger.info("Models cleared.")

# ── FastAPI app ────────────────────────────────────────────────────────────────
app = FastAPI(
    title="CNN Latency Experiment API",
    description="Server inferensi untuk pengukuran latensi HTTP endpoint — paper komparatif CNN penyakit daun tomat",
    version="1.0.0",
    lifespan=lifespan,
)

# ── Health check ───────────────────────────────────────────────────────────────
@app.get("/health")
def health():
    return {
        "status": "ok",
        "loaded_models": list(LOADED_MODELS.keys()),
        "total_loaded": len(LOADED_MODELS),
    }

# ── Debug endpoint ─────────────────────────────────────────────────────────────
@app.get("/debug")
def debug():
    """Diagnostik: cek apakah file model ada dan ukurannya benar."""
    model_files = {}
    for name, config in MODEL_CONFIG.items():
        path = config["path"]
        exists = os.path.exists(path)
        size = os.path.getsize(path) if exists else 0
        # Cek apakah file LFS pointer (< 1KB = pointer, bukan model asli)
        is_lfs_pointer = False
        if exists and size < 1024:
            with open(path, "rb") as f:
                head = f.read(100)
            is_lfs_pointer = b"version https://git-lfs" in head
        model_files[name] = {
            "path": path,
            "exists": exists,
            "size_bytes": size,
            "size_mb": round(size / 1024 / 1024, 2) if exists else 0,
            "is_lfs_pointer": is_lfs_pointer,
        }
    return {
        "model_dir": MODEL_DIR,
        "model_dir_exists": os.path.isdir(MODEL_DIR),
        "files": model_files,
        "loaded": list(LOADED_MODELS.keys()),
        "load_errors": LOAD_ERRORS,
        "tf_version": tf.__version__,
        "keras_version": keras.__version__,
    }

# ── List model yang tersedia ───────────────────────────────────────────────────
@app.get("/models")
def list_models():
    return {
        "available": list(LOADED_MODELS.keys()),
        "all_configured": list(MODEL_CONFIG.keys()),
    }

# ── Endpoint predict utama ─────────────────────────────────────────────────────
@app.post("/predict/{model_name}")
async def predict(
    model_name: str,
    file: UploadFile = File(...),
):
    """
    Endpoint inferensi per model.

    Args:
        model_name : salah satu dari 7 nama model (lihat MODEL_CONFIG)
        file       : gambar daun tomat (JPG/PNG)

    Returns:
        JSON berisi prediksi, confidence, dan breakdown latensi (ms)
    """
    # Validasi nama model
    if model_name not in MODEL_CONFIG:
        raise HTTPException(
            status_code=404,
            detail=f"Model '{model_name}' tidak dikenal. "
                   f"Pilihan: {list(MODEL_CONFIG.keys())}"
        )

    if model_name not in LOADED_MODELS:
        raise HTTPException(
            status_code=503,
            detail=f"Model '{model_name}' tidak berhasil dimuat saat startup."
        )

    model = LOADED_MODELS[model_name]
    preprocess_mode = MODEL_CONFIG[model_name]["preprocess"]

    # ── Baca dan decode gambar ─────────────────────────────
    t_read_start = time.perf_counter()
    try:
        img_bytes = await file.read()
        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
        img = img.resize(IMAGE_SIZE, Image.BILINEAR)
        img_array = np.array(img, dtype=np.float32)[np.newaxis, ...]  # (1,224,224,3)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Gagal membaca gambar: {e}")
    t_read_end = time.perf_counter()

    # ── Preprocess ─────────────────────────────────────────
    t_pre_start = time.perf_counter()
    img_array = preprocess_image(img_array, preprocess_mode)
    t_pre_end = time.perf_counter()

    # ── Inferensi ──────────────────────────────────────────
    t_infer_start = time.perf_counter()
    preds = model.predict(img_array, verbose=0)
    t_infer_end = time.perf_counter()

    # ── Hasil ──────────────────────────────────────────────
    class_idx  = int(np.argmax(preds[0]))
    confidence = float(np.max(preds[0]))

    # Breakdown latensi dalam ms
    read_ms   = round((t_read_end   - t_read_start)  * 1000, 3)
    pre_ms    = round((t_pre_end    - t_pre_start)   * 1000, 3)
    infer_ms  = round((t_infer_end  - t_infer_start) * 1000, 3)
    total_ms  = round(read_ms + pre_ms + infer_ms, 3)

    return JSONResponse(content={
        "model":       model_name,
        "prediction":  CLASS_NAMES[class_idx],
        "class_index": class_idx,
        "confidence":  round(confidence, 6),
        "latency": {
            "read_decode_ms":  read_ms,
            "preprocess_ms":   pre_ms,
            "inference_ms":    infer_ms,
            "total_server_ms": total_ms,
        },
    })