Update app.py

app.py

@@ -3,69 +3,39 @@ import os
 import time
 import numpy as np
 import cv2
-import glob
-
-# Deep Learning Libraries
-import tensorflow as tf
-from tensorflow.keras.models import load_model, Model
-from tensorflow.keras.preprocessing.image import load_img, img_to_array, ImageDataGenerator
-from tensorflow.keras.applications import EfficientNetB0
-from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, Dropout, Input
-from tensorflow.keras.optimizers import Adam
-from tensorflow.keras.callbacks import ModelCheckpoint, ReduceLROnPlateau
-from sklearn.utils import class_weight
+
+# Deep Learning Libraries (PyTorch)
+import torch
+import torch.nn as nn
+from torchvision import models, transforms
 
 app = Flask(__name__)
 
 # -------------------- CONFIG --------------------
 UPLOAD_FOLDER = "uploads"
-DATASET_PATH = os.path.join("static", "dataset", "colored_images")
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 
-# 5 CLASSES (Correct Order)
-CLASSES = ["No_DR", "Mild", "Moderate", "Severe", "Proliferate_DR"]
+# Important: Update these paths to where your .pth files are saved locally!
+DENSENET_PATH = "dr_model_final.pth"
+EFFICIENTNET_PATH = "efficientnet_dr_model.pth"
 
-MODEL_FILE = "retina_efficientnet_v2.h5" 
-IMG_SIZE = (224, 224)
-BATCH_SIZE = 16
+# Exact class order used by the PyTorch ImageFolder during training
+PYTORCH_CLASSES = ['Mild', 'Moderate', 'No_DR', 'Proliferative', 'Severe']
+NUM_CLASSES = len(PYTORCH_CLASSES)
 
-os.makedirs(UPLOAD_FOLDER, exist_ok=True)
-MODEL = None
-LOAD_ERROR = None
-
-# -------------------- 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
-
-from tensorflow.keras.layers import (
-    Conv2D, BatchNormalization, Activation, DepthwiseConv2D, 
-    Rescaling, ZeroPadding2D, Add, Multiply, InputLayer
-)
-
-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)
-}
-
-# -------------------- ADVANCED PREPROCESSING --------------------
+DENSENET_MODEL = None
+EFFICIENTNET_MODEL = None
+DEVICE = None
+
+# -------------------- PREPROCESSING --------------------
 def enhance_medical_image(img):
+    """Applies CLAHE enhancement to retina images."""
     try:
-        if np.max(img) <= 1.0: img = (img * 255).astype(np.uint8)
-        else: img = img.astype(np.uint8)
+        # Convert to uint8 if not already
+        if np.max(img) <= 1.0: 
+            img = (img * 255).astype(np.uint8)
+        else: 
+            img = img.astype(np.uint8)
 
         lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
         l, a, b = cv2.split(lab)
@@ -74,12 +44,13 @@ def enhance_medical_image(img):
         limg = cv2.merge((cl, a, b))
         final = cv2.cvtColor(limg, cv2.COLOR_LAB2RGB)
         
-        return final.astype(np.float32) / 255.0
+        # Return uint8 array for PyTorch ToPILImage transform
+        return final
     except Exception:
-        return img.astype(np.float32) / 255.0
+        return img.astype(np.uint8)
 
-# -------------------- EYE VALIDATION --------------------
 def is_valid_eye_image(img_path):
+    """Validation to ensure uploaded images are likely retina scans."""
     try:
         img = cv2.imread(img_path)
         if img is None: return False
@@ -98,159 +69,111 @@ def is_valid_eye_image(img_path):
         if mean_brightness < 5 or mean_brightness > 240: return False
 
         return True
-    except Exception: return True
-
-# -------------------- TRAINING PIPELINE --------------------
-def build_model(trainable=False):
-    base_model = EfficientNetB0(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
-    base_model.trainable = trainable
-
-    x = base_model.output
-    x = GlobalAveragePooling2D()(x)
-    x = Dense(512, activation='relu')(x)
-    x = Dropout(0.5)(x) 
-    predictions = Dense(len(CLASSES), activation='softmax')(x)
-
-    model = Model(inputs=base_model.input, outputs=predictions)
-    return model
-
-def train_on_dataset():
-    print(f"\n[TRAINING] 🚀 Starting 5-Class Training Pipeline (Fixed 35 Epochs)...")
-    
-    if not os.path.exists(DATASET_PATH):
-        print("[ERROR] Dataset not found.")
-        return None
-
-    train_datagen = ImageDataGenerator(
-        preprocessing_function=enhance_medical_image,
-        rotation_range=20,
-        width_shift_range=0.1,
-        height_shift_range=0.1,
-        zoom_range=0.1,
-        horizontal_flip=True,
-        vertical_flip=True,
-        fill_mode='nearest',
-        validation_split=0.2
-    )
-
-    train_generator = train_datagen.flow_from_directory(
-        DATASET_PATH, target_size=IMG_SIZE, batch_size=BATCH_SIZE,
-        class_mode='categorical', subset='training',
-        classes=CLASSES, shuffle=True
-    )
-
-    val_generator = train_datagen.flow_from_directory(
-        DATASET_PATH, target_size=IMG_SIZE, batch_size=BATCH_SIZE,
-        class_mode='categorical', subset='validation',
-        classes=CLASSES, shuffle=False
-    )
+    except Exception: 
+        return True
 
+def calculate_entropy(img_array):
     try:
-        train_classes = train_generator.classes
-        class_weights = class_weight.compute_class_weight(
-            class_weight='balanced',
-            classes=np.unique(train_classes),
-            y=train_classes
-        )
-        class_weights_dict = dict(enumerate(class_weights))
-        print(f"[INFO] Class Weights: {class_weights_dict}")
-    except:
-        class_weights_dict = None
-
-    # STAGE 1
-    print("\n[STAGE 1] Warming up Head...")
-    model = build_model(trainable=False)
-    model.compile(optimizer=Adam(learning_rate=1e-3), loss='categorical_crossentropy', metrics=['accuracy'])
-    
-    model.fit(
-        train_generator,
-        validation_data=val_generator,
-        epochs=5, 
-        class_weight=class_weights_dict,
-        verbose=1
+        gray = cv2.cvtColor(img_array, 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]
+        return float(-np.sum(hist_norm * np.log2(hist_norm)))
+    except: 
+        return 4.5
+
+# -------------------- MODEL BUILDERS --------------------
+def build_densenet():
+    model = models.densenet121(weights=None)
+    num_ftrs = model.classifier.in_features
+    model.classifier = nn.Sequential(
+        nn.Linear(num_ftrs, 512),
+        nn.ReLU(),
+        nn.Dropout(0.4), 
+        nn.Linear(512, NUM_CLASSES)
     )
+    return model
 
-    # STAGE 2
-    print("\n[STAGE 2] Unfreezing ENTIRE model (Training for 35 Epochs)...")
-    
-    for layer in model.layers:
-        layer.trainable = True
-    
-    model.compile(optimizer=Adam(learning_rate=1e-4), loss='categorical_crossentropy', metrics=['accuracy'])
-
-    # REMOVED EarlyStopping to force full 35 epochs
-    callbacks = [
-        ModelCheckpoint(MODEL_FILE, monitor='val_accuracy', save_best_only=True, verbose=1),
-        ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=5, min_lr=1e-7, verbose=1)
-    ]
-
-    model.fit(
-        train_generator,
-        validation_data=val_generator,
-        epochs=35, # Fixed at 35 epochs
-        class_weight=class_weights_dict,
-        callbacks=callbacks,
-        verbose=1
+def build_efficientnet():
+    model = models.efficientnet_b4(weights=None)
+    num_ftrs = model.classifier[1].in_features
+    model.classifier = nn.Sequential(
+        nn.Dropout(p=0.5, inplace=True), 
+        nn.Linear(num_ftrs, 512),
+        nn.BatchNorm1d(512), 
+        nn.ReLU(),
+        nn.Dropout(p=0.5), 
+        nn.Linear(512, NUM_CLASSES)
     )
-
-    print(f"[SUCCESS] Training complete.")
-    
-    # RELOAD BEST MODEL
-    # Since we removed EarlyStopping, the model in memory is the LAST epoch.
-    # We reload the file to ensure we use the BEST epoch saved by ModelCheckpoint.
-    print(f"[INFO] Reloading best model from {MODEL_FILE}...")
-    try:
-        model = load_model(MODEL_FILE, custom_objects=CUSTOM_OBJECTS)
-    except Exception as e:
-        print(f"[WARN] Could not reload best model ({e}). Using last epoch weights.")
-
     return model
 
 # -------------------- LOAD / INIT --------------------
-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}")
-            MODEL = train_on_dataset()
+def init_models():
+    global DENSENET_MODEL, EFFICIENTNET_MODEL, DEVICE
+    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"[INIT] Initializing Ensemble on {DEVICE}...")
+
+    # Load DenseNet
+    DENSENET_MODEL = build_densenet()
+    if os.path.exists(DENSENET_PATH):
+        checkpoint = torch.load(DENSENET_PATH, map_location=DEVICE)
+        DENSENET_MODEL.load_state_dict(checkpoint['model_state_dict'])
+        print(f"[INIT] DenseNet loaded successfully.")
     else:
-        print(f"[INIT] No model found. Starting training...")
-        MODEL = train_on_dataset()
-
-# -------------------- INFERENCE HELPERS --------------------
-def process_single_image(image_path):
-    try:
-        img = load_img(image_path, target_size=IMG_SIZE)
-        img_array = img_to_array(img)
-        img_enhanced = enhance_medical_image(img_array)
-        return np.expand_dims(img_enhanced, axis=0)
-    except Exception as e: return str(e)
-
-def calculate_entropy(img_array):
-    try:
-        if img_array.dtype != np.uint8:
-            calc_img = (img_array * 255).astype(np.uint8) if np.max(img_array) <= 1.0 else 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]
-        return float(-np.sum(hist_norm * np.log2(hist_norm)))
-    except: return 4.5
-
-def predict_smart(model, input_batch):
-    img = input_batch[0]
-    aug_batch = np.array([img, np.fliplr(img), np.flipud(img)])
-    preds = model.predict(aug_batch)
-    avg_pred = np.mean(preds, axis=0)
-    return avg_pred
+        print(f"[WARNING] DenseNet file missing at {DENSENET_PATH}. App will crash on inference.")
+    DENSENET_MODEL = DENSENET_MODEL.to(DEVICE)
+    DENSENET_MODEL.eval()
+
+    # Load EfficientNet
+    EFFICIENTNET_MODEL = build_efficientnet()
+    if os.path.exists(EFFICIENTNET_PATH):
+        checkpoint = torch.load(EFFICIENTNET_PATH, map_location=DEVICE)
+        EFFICIENTNET_MODEL.load_state_dict(checkpoint['model_state_dict'])
+        print(f"[INIT] EfficientNet loaded successfully.")
+    else:
+        print(f"[WARNING] EfficientNet file missing at {EFFICIENTNET_PATH}. App will crash on inference.")
+    EFFICIENTNET_MODEL = EFFICIENTNET_MODEL.to(DEVICE)
+    EFFICIENTNET_MODEL.eval()
+
+# -------------------- ENSEMBLE INFERENCE --------------------
+def predict_ensemble(img_array):
+    # PyTorch transforms (requires uint8 array as input to ToPILImage)
+    transform_dense = transforms.Compose([
+        transforms.ToPILImage(),
+        transforms.Resize((256, 256)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ])
+    
+    transform_eff = transforms.Compose([
+        transforms.ToPILImage(),
+        transforms.Resize((288, 288)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ])
+
+    # Convert directly to device tensors
+    t_dense = transform_dense(img_array).unsqueeze(0).to(DEVICE)
+    t_eff = transform_eff(img_array).unsqueeze(0).to(DEVICE)
+
+    with torch.no_grad():
+        # --- DenseNet TTA ---
+        d_out1 = torch.softmax(DENSENET_MODEL(t_dense), dim=1)
+        d_out2 = torch.softmax(DENSENET_MODEL(torch.flip(t_dense, dims=[3])), dim=1) # Horiz flip
+        d_out3 = torch.softmax(DENSENET_MODEL(torch.flip(t_dense, dims=[2])), dim=1) # Vert flip
+        d_probs = (d_out1 + d_out2 + d_out3) / 3.0
+
+        # --- EfficientNet TTA ---
+        e_out1 = torch.softmax(EFFICIENTNET_MODEL(t_eff), dim=1)
+        e_out2 = torch.softmax(EFFICIENTNET_MODEL(torch.flip(t_eff, dims=[3])), dim=1)
+        e_out3 = torch.softmax(EFFICIENTNET_MODEL(torch.flip(t_eff, dims=[2])), dim=1)
+        e_probs = (e_out1 + e_out2 + e_out3) / 3.0
+
+        # --- Soft Voting Ensemble ---
+        ensemble_probs = (d_probs + e_probs) / 2.0
+        confidence, pred_idx = torch.max(ensemble_probs, 1)
+
+        return PYTORCH_CLASSES[pred_idx.item()], confidence.item()
 
 # -------------------- ROUTES --------------------
 @app.route("/")
@@ -259,7 +182,9 @@ def index():
 
 @app.route("/analyze", methods=["POST"])
 def analyze():
-    if "image" not in request.files: return jsonify({"error": "No image"}), 400
+    if "image" not in request.files: 
+        return jsonify({"error": "No image uploaded"}), 400
+        
     file = request.files["image"]
     temp_path = os.path.join(UPLOAD_FOLDER, f"scan_{int(time.time())}.jpg")
     file.save(temp_path)
@@ -267,53 +192,49 @@ def analyze():
     try:
         if not is_valid_eye_image(temp_path):
             return jsonify({
-                "diagnosis": "Invalid Image",
-                "severity": "Scan Rejected",
-                "description": "The image detected does not appear to be a standard retina scan. Please ensure proper alignment.",
+                "diagnosis": "Invalid Image", "severity": "Scan Rejected",
+                "description": "The image detected does not appear to be a standard retina scan.",
                 "confidence": "0%", "color": "rose", "icon": "alert-circle", "features": {"entropy": "N/A"}
             })
 
-        if MODEL is None:
-            return jsonify({
-                "diagnosis": "System Error",
-                "description": "Model is training.",
-                "confidence": "0%", "color": "rose", "icon": "alert-octagon"
-            })
-
-        input_data = process_single_image(temp_path)
-        if isinstance(input_data, str):
-            return jsonify({"diagnosis": "Error", "description": input_data, "confidence": "0%", "color": "rose", "icon": "alert-triangle"})
+        # Load image via OpenCV and apply enhancements
+        img = cv2.imread(temp_path)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img_enhanced = enhance_medical_image(img)
+        
+        # Calculate Entropy
+        ent = calculate_entropy(img_enhanced)
 
-        preds = predict_smart(MODEL, input_data)
-        idx = np.argmax(preds)
-        label = CLASSES[idx]
-        conf = preds[idx] * 100
-        ent = calculate_entropy(input_data)
+        # Run Ensemble Prediction
+        label, conf = predict_ensemble(img_enhanced)
+        conf_percentage = conf * 100
 
-        # Updated Mapping for 5 Classes
+        # Map PyTorch labels to Front-end UI UI mapping
         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"),
-            "Proliferate_DR": ("Proliferative DR", "Stage 4", "purple", "alert-octagon"),
+            "Proliferative": ("Proliferative DR", "Stage 4", "purple", "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: {diag}",
-            "confidence": f"{conf:.1f}%",
+            "description": f"Ensemble AI Analysis: {diag}",
+            "confidence": f"{conf_percentage:.1f}%",
             "features": {"entropy": f"{ent:.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)
+        if os.path.exists(temp_path): 
+            os.remove(temp_path)
 
 # -------------------- MAIN --------------------
-init_model()
+init_models()
 
 if __name__ == "__main__":
     app.run(debug=True, port=7860)