main.py

import os

os.environ["TF_USE_LEGACY_KERAS"] = "1"

import io
import logging
import numpy as np
import pandas as pd
import cv2
from pathlib import Path
from contextlib import asynccontextmanager

import uvicorn
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
from sklearn.preprocessing import LabelEncoder
import tensorflow as tf

# Use tf_keras (Keras 2) to stay compatible with the saved .h5 model
try:
    import tf_keras as keras
    from tf_keras.layers import (
        Input, Conv2D, BatchNormalization, Activation,
        MaxPooling2D, Add, GlobalAveragePooling2D, Dense, Dropout
    )
    from tf_keras.models import Model
    from tf_keras import regularizers
    from tf_keras.optimizers import AdamW
    logger_init = "tf_keras (Keras 2) loaded ✓"
except ImportError:
    # Fallback: standard tf.keras
    keras = tf.keras
    from tensorflow.keras.layers import (
        Input, Conv2D, BatchNormalization, Activation,
        MaxPooling2D, Add, GlobalAveragePooling2D, Dense, Dropout
    )
    from tensorflow.keras.models import Model
    from tensorflow.keras import regularizers
    from tensorflow.keras.optimizers import AdamW
    logger_init = "tf.keras (fallback) loaded ✓"

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logger.info(logger_init)

MODEL_PATH = os.getenv("MODEL_PATH", "IrisRecognizer95.h5")
CSV_PATH   = os.getenv("CSV_PATH",   "Banned_travelers22.csv")

IMG_HEIGHT   = 150
IMG_WIDTH    = 150
NUM_CHANNELS = 1
NUM_CLASSES  = 2000

model:         object            = None
banned_df:     pd.DataFrame      = None
label_encoder: LabelEncoder      = None


#  Architecture (identical to training notebook) 
def residual_block(x, filters, kernel_size=3, stride=1):
    shortcut = x
    x = Conv2D(filters, kernel_size, strides=stride, padding="same",
               kernel_regularizer=regularizers.l2(1e-4))(x)
    x = BatchNormalization()(x)
    x = Activation("relu")(x)
    x = Conv2D(filters, kernel_size, padding="same",
               kernel_regularizer=regularizers.l2(1e-4))(x)
    x = BatchNormalization()(x)
    if stride != 1 or shortcut.shape[-1] != filters:
        shortcut = Conv2D(filters, 1, strides=stride, padding="same")(shortcut)
        shortcut = BatchNormalization()(shortcut)
    x = Add()([x, shortcut])
    x = Activation("relu")(x)
    return x


def build_model(input_shape=(IMG_HEIGHT, IMG_WIDTH, NUM_CHANNELS),
                num_classes=NUM_CLASSES):
    inputs = Input(shape=input_shape)
    x = Conv2D(64, 7, strides=2, padding="same")(inputs)
    x = BatchNormalization()(x)
    x = Activation("relu")(x)
    x = MaxPooling2D(3, strides=2, padding="same")(x)
    x = residual_block(x, 64);  x = residual_block(x, 64)
    x = residual_block(x, 128, stride=2); x = residual_block(x, 128)
    x = residual_block(x, 256, stride=2); x = residual_block(x, 256)
    x = residual_block(x, 512, stride=2); x = residual_block(x, 512)
    x = GlobalAveragePooling2D()(x)
    x = Dropout(0.4)(x)
    outputs = Dense(num_classes, activation="softmax")(x)
    m = Model(inputs, outputs, name="IRIS_ResNet")
    m.compile(optimizer=AdamW(learning_rate=0.0003, weight_decay=1e-4),
              loss="sparse_categorical_crossentropy", metrics=["accuracy"])
    return m


def load_model_safe(model_path: str):
    """
    Try standard load first.
    On Keras version mismatch → rebuild architecture + load weights only.
    """
    # 1. Try with tf_keras directly
    try:
        logger.info("Attempt 1: tf_keras.models.load_model …")
        m = keras.models.load_model(model_path)
        logger.info("Load succeeded ✓")
        return m
    except Exception as e:
        logger.warning(f"Attempt 1 failed: {e}")

    # 2. Try with tf.keras
    try:
        logger.info("Attempt 2: tf.keras.models.load_model …")
        m = tf.keras.models.load_model(model_path)
        logger.info("Load succeeded ✓")
        return m
    except Exception as e:
        logger.warning(f"Attempt 2 failed: {e}")

    # 3. Rebuild + weights only
    logger.info("Attempt 3: rebuild architecture + load_weights …")
    m = build_model()
    m.load_weights(model_path)
    logger.info("Weights-only load succeeded ✓")
    return m


# ─ Preprocessing 
def resize_keep_aspect_ratio(img, target_h=IMG_HEIGHT, target_w=IMG_WIDTH, pad_value=255):
    aspect = img.shape[1] / img.shape[0]
    if aspect > target_w / target_h:
        new_w, new_h = target_w, int(target_w / aspect)
    else:
        new_h, new_w = target_h, int(target_h * aspect)
    resized = cv2.resize(img, (new_w, new_h))
    padded  = np.full((target_h, target_w), pad_value, dtype=np.uint8)
    padded[(target_h - new_h) // 2:(target_h - new_h) // 2 + new_h,
           (target_w - new_w) // 2:(target_w - new_w) // 2 + new_w] = resized
    return padded


def preprocess_image_bytes(image_bytes: bytes) -> np.ndarray:
    img_np = np.array(Image.open(io.BytesIO(image_bytes)).convert("L"), dtype=np.uint8)
    img    = resize_keep_aspect_ratio(img_np).astype(np.float32) / 255.0
    return img.reshape(1, IMG_HEIGHT, IMG_WIDTH, NUM_CHANNELS)


#  Lifespan 
@asynccontextmanager
async def lifespan(app: FastAPI):
    global model, banned_df, label_encoder

    if not Path(MODEL_PATH).exists():
        raise FileNotFoundError(f"Model not found: {MODEL_PATH}")

    logger.info(f"Loading model from {MODEL_PATH} …")
    model = load_model_safe(MODEL_PATH)
    logger.info("Model ready ✓")

    if not Path(CSV_PATH).exists():
        raise FileNotFoundError(f"CSV not found: {CSV_PATH}")

    banned_df = pd.read_csv(CSV_PATH)
    banned_df.columns = banned_df.columns.str.strip()
    for col in ("Label", "person_id", "Status"):
        banned_df[col] = banned_df[col].astype(str).str.strip()

    label_encoder = LabelEncoder()
    label_encoder.fit(banned_df["Label"].unique())
    logger.info(f"Ready — {len(label_encoder.classes_)} classes | {len(banned_df)} CSV rows ")

    yield
    logger.info("Shutdown.")


#  App 
app = FastAPI(
    title="Iris Recognition — Banned Traveler Detection",
    description="Upload an iris image → get Banned / Allowed status.",
    version="1.0.0",
    lifespan=lifespan,
)

app.add_middleware(CORSMiddleware, allow_origins=["*"],
                   allow_methods=["*"], allow_headers=["*"])


@app.get("/")
def root():
    return {"message": "Iris API running ", "docs": "/docs", "health": "/health"}


@app.get("/health")
def health():
    return {
        "status":         "ok",
        "model_loaded":   model is not None,
        "csv_loaded":     banned_df is not None,
        "num_classes":    int(len(label_encoder.classes_)) if label_encoder else 0,
        "banned_records": int(len(banned_df)) if banned_df is not None else 0,
    }


@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    """Upload an iris image → returns person_id, status (Banned/Allowed), confidence."""
    if file.content_type not in ("image/jpeg", "image/png", "image/jpg", "image/bmp"):
        raise HTTPException(415, f"Unsupported format: {file.content_type}")

    image_bytes = await file.read()
    if not image_bytes:
        raise HTTPException(400, "Empty file.")

    try:
        img_array = preprocess_image_bytes(image_bytes)
    except Exception as e:
        raise HTTPException(422, f"Image processing error: {e}")

    probs           = model.predict(img_array, verbose=0)
    pred_idx        = int(np.argmax(probs[0]))
    confidence      = float(np.max(probs[0]))
    predicted_label = str(label_encoder.classes_[pred_idx])

    match = banned_df[banned_df["Label"] == predicted_label]
    if not match.empty:
        person_id = str(match.iloc[0]["person_id"])
        status    = str(match.iloc[0]["Status"])
    else:
        person_id = predicted_label.split("-")[0]
        status    = "Unknown"

    is_banned = status.lower() == "banned"
    logger.info(f"[PREDICT] {predicted_label} | {status} | conf={confidence:.4f}")

    return JSONResponse({
        "person_id":       person_id,
        "predicted_label": predicted_label,
        "status":          status,
        "confidence":      round(confidence, 4),
        "is_banned":       is_banned,
        "message": (
            f" BANNED — Person {person_id} is NOT allowed to travel." if is_banned else
            f" ✓ ALLOWED — Person {person_id} is cleared to travel."    if status == "Allowed" else
            f" UNKNOWN — Person {person_id} not found in records."
        ),
    })


if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=False)