app.py

from flask import Flask, render_template, request, jsonify
import os
import time
import numpy as np
import cv2

# Deep Learning Libraries
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import load_img, img_to_array

# Import all layers used in EfficientNet and our custom head to patch them
from tensorflow.keras.layers import (
    Dense, GlobalAveragePooling2D, Dropout, Conv2D, BatchNormalization, 
    Activation, DepthwiseConv2D, Rescaling, ZeroPadding2D, Add, Multiply, InputLayer
)

app = Flask(__name__)

# -------------------- CONFIG --------------------
UPLOAD_FOLDER = "uploads"
CLASSES = ["No_DR", "Mild", "Moderate", "Severe"]
MODEL_FILE = "retina_efficientnet_v2.h5" 
IMG_SIZE = (224, 224)

os.makedirs(UPLOAD_FOLDER, exist_ok=True)
MODEL = None
LOAD_ERROR = None  # Store the specific reason for failure

# -------------------- COMPATIBILITY FIX --------------------
def fix_layer_config(cls):
    class FixedLayer(cls):
        def __init__(self, *args, **kwargs):
            kwargs.pop('quantization_config', None)
            kwargs.pop('glitch_filter', None) 
            super().__init__(*args, **kwargs)
    return FixedLayer

CUSTOM_OBJECTS = {
    'Dense': fix_layer_config(Dense),
    'Dropout': fix_layer_config(Dropout),
    'GlobalAveragePooling2D': fix_layer_config(GlobalAveragePooling2D),
    'Conv2D': fix_layer_config(Conv2D),
    'BatchNormalization': fix_layer_config(BatchNormalization),
    'Activation': fix_layer_config(Activation),
    'DepthwiseConv2D': fix_layer_config(DepthwiseConv2D),
    'Rescaling': fix_layer_config(Rescaling),
    'ZeroPadding2D': fix_layer_config(ZeroPadding2D),
    'Add': fix_layer_config(Add),
    'Multiply': fix_layer_config(Multiply),
    'InputLayer': fix_layer_config(InputLayer)
}

# -------------------- LOAD MODEL --------------------
def init_model():
    global MODEL, LOAD_ERROR
    LOAD_ERROR = None
    
    if os.path.exists(MODEL_FILE):
        print(f"[INIT] Model found: {MODEL_FILE}")
        try:
            MODEL = load_model(MODEL_FILE, custom_objects=CUSTOM_OBJECTS)
            print("[INIT] Model loaded successfully.")
        except Exception as e:
            print(f"[ERROR] Failed to load model: {e}")
            LOAD_ERROR = str(e)
            MODEL = None
    else:
        print(f"[ERROR] Model file '{MODEL_FILE}' NOT FOUND on server.")
        MODEL = None

# -------------------- PREPROCESSING & TTA --------------------
def calculate_entropy(img_array):
    try:
        if img_array.dtype != np.uint8:
            if np.max(img_array) <= 1.0: calc_img = (img_array * 255).astype(np.uint8)
            else: calc_img = img_array.astype(np.uint8)
        else: calc_img = img_array

        if len(calc_img.shape) == 4: calc_img = calc_img[0]
        gray = cv2.cvtColor(calc_img, cv2.COLOR_RGB2GRAY)
        hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
        hist_norm = hist.ravel() / hist.sum()
        hist_norm = hist_norm[hist_norm > 0]
        entropy_val = -np.sum(hist_norm * np.log2(hist_norm))
        return float(entropy_val)
    except Exception as e:
        return 4.5

def process_single_image(image_path):
    try:
        img = load_img(image_path, target_size=IMG_SIZE)
        img_array = img_to_array(img)
        
        # CLAHE Logic
        if np.max(img_array) <= 1.0: img_array = (img_array * 255).astype(np.uint8)
        else: img_array = img_array.astype(np.uint8)

        lab = cv2.cvtColor(img_array, cv2.COLOR_RGB2LAB)
        l, a, b = cv2.split(lab)
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
        cl = clahe.apply(l)
        limg = cv2.merge((cl, a, b))
        final = cv2.cvtColor(limg, cv2.COLOR_LAB2RGB)
        
        return np.expand_dims(final.astype(np.float32) / 255.0, axis=0)
    except Exception as e:
        return str(e)

def predict_with_tta(model, input_batch):
    """
    Test Time Augmentation:
    Predicts on the original image + flipped versions and averages the results.
    """
    img = input_batch[0] # Extract image from batch (224, 224, 3)
    
    # Create batch of 3 variants: Original, Horizontal Flip, Vertical Flip
    # Retina images have no "correct" up/down, so vertical flipping is valid logic.
    aug_batch = np.array([
        img,
        np.fliplr(img),
        np.flipud(img)
    ])
    
    # Get predictions for all 3 variations
    preds = model.predict(aug_batch)
    
    # Average the probabilities across the 3 views
    avg_pred = np.mean(preds, axis=0)
    return avg_pred

# -------------------- ROUTES --------------------
@app.route("/")
def index():
    return render_template("index.html")

@app.route("/analyze", methods=["POST"])
def analyze():
    if "image" not in request.files: return jsonify({"error": "No image"}), 400
    file = request.files["image"]
    temp_path = os.path.join(UPLOAD_FOLDER, f"scan_{int(time.time())}.jpg")
    file.save(temp_path)

    try:
        # 1. Check if Model Loaded
        if MODEL is None:
            if not os.path.exists(MODEL_FILE):
                return jsonify({
                    "diagnosis": "System Error",
                    "description": f"Model file '{MODEL_FILE}' not found.",
                    "confidence": "0%",
                    "color": "rose", "icon": "alert-octagon"
                })
            else:
                error_msg = LOAD_ERROR if LOAD_ERROR else "Model initialization skipped by server."
                return jsonify({
                    "diagnosis": "Load Error",
                    "description": f"Error: {error_msg}", 
                    "confidence": "0%",
                    "color": "rose", "icon": "alert-octagon"
                })

        # 2. Process Image
        input_data = process_single_image(temp_path)
        if isinstance(input_data, str):
            return jsonify({
                "diagnosis": "OpenCV Error",
                "description": f"Processing failed: {input_data}",
                "confidence": "0%",
                "color": "rose", "icon": "alert-triangle"
            })

        # 3. Predict with TTA (Smart Averaging)
        preds = predict_with_tta(MODEL, input_data)
        
        idx = np.argmax(preds)
        label = CLASSES[idx]
        conf = preds[idx] * 100
        
        # 4. Calculate Real Entropy
        entropy_val = calculate_entropy(input_data)

        # FORMAT RESULT
        mapping = {
            "No_DR": ("No DR", "Normal", "emerald", "check-circle"),
            "Mild": ("Mild DR", "Stage 1", "yellow", "alert-triangle"),
            "Moderate": ("Moderate DR", "Stage 2", "orange", "alert-triangle"),
            "Severe": ("Severe DR", "Stage 3", "rose", "alert-octagon"),
        }
        diag, sev, col, icon = mapping.get(label, ("Unknown", "-", "gray", "help-circle"))

        return jsonify({
            "diagnosis": diag, "severity": sev, "color": col, "icon": icon,
            "description": f"AI Analysis Result: {diag}",
            "confidence": f"{conf:.1f}%",
            "features": {"entropy": f"{entropy_val:.3f}"}
        })

    except Exception as e:
        return jsonify({
            "diagnosis": "Crash",
            "description": str(e),
            "confidence": "0%",
            "color": "rose", "icon": "x-octagon"
        })
    finally:
        if os.path.exists(temp_path): os.remove(temp_path)

# -------------------- INITIALIZE ON IMPORT --------------------
init_model()

if __name__ == "__main__":
    app.run(debug=True, port=7860)