Update app.py

app.py

@@ -1,10 +1,11 @@
 from flask import Flask, request, jsonify, send_file
-from tensorflow.keras.models import load_model, Model
 from PIL import Image
-import numpy as np
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
 import os
 import cv2
-import tensorflow as tf
+import numpy as np
 from datetime import datetime
 import sqlite3
 
@@ -17,6 +18,7 @@ if not os.path.exists(OUTPUT_DIR):
 
 DB_PATH = os.path.join(OUTPUT_DIR, 'results.db')
 
+
 def init_db():
     conn = sqlite3.connect(DB_PATH)
     cursor = conn.cursor()
@@ -33,65 +35,27 @@ def init_db():
     conn.commit()
     conn.close()
 
+
 init_db()
 
-# ✅ Load Model
-model = load_model('mobilenet_glaucoma_model.h5', compile=False)
+# ✅ Load PyTorch Model
+model = torch.load('your_model.pt', map_location=torch.device('cpu'))
+model.eval()
 
 # ✅ Preprocess Image
-def preprocess_image(img):
-    img = img.resize((224, 224))
-    img = np.array(img) / 255.0
-    img = np.expand_dims(img, axis=0)
-    return img
-
-# ✅ Grad-CAM Generation
-def make_gradcam(img_array, model, last_conv_layer_name='Conv_1_bn'):
-    """Generate Grad-CAM for the given image and model."""
-    last_conv_layer = model.get_layer(last_conv_layer_name)
-    grad_model = Model(inputs=model.inputs, outputs=[last_conv_layer.output, model.output])
-
-    with tf.GradientTape() as tape:
-        conv_outputs, predictions = grad_model(img_array)
-        loss = predictions[:, 0]
-    grads = tape.gradient(loss, conv_outputs)
-
-    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
-    conv_outputs = conv_outputs[0].numpy()
-    pooled_grads = pooled_grads.numpy()
-
-    for i in range(conv_outputs.shape[-1]):
-        conv_outputs[..., i] *= pooled_grads[i]
-
-    heatmap = tf.reduce_mean(conv_outputs, axis=-1).numpy()
-    heatmap = np.maximum(heatmap, 0)
-    heatmap /= np.max(heatmap)
-
-    return heatmap
-
-# ✅ Save Grad-CAM Overlay
-def save_gradcam_image(original_img, heatmap, filename='gradcam.png', output_dir=OUTPUT_DIR):
-    """Save the Grad-CAM overlay image and return its path."""
-    img = np.array(original_img.resize((224, 224)))
-    heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]))
-    heatmap = np.uint8(255 * heatmap)
-
-    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
-    overlay = cv2.addWeighted(img, 0.6, heatmap, 0.4, 0)
-
-    filepath = os.path.join(output_dir, filename)
-    cv2.imwrite(filepath, overlay)
-
-    return filepath
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+])
 
 @app.route('/')
 def home():
-    return "Glaucoma Detection Flask API is running!"
+    return "Glaucoma Detection Flask API (PyTorch) is running!"
 
 @app.route("/test_file")
 def test_file():
-    """Check if the model file is present and readable."""
-    filepath = "mobilenet_glaucoma_model.h5"
+    """Check if the .pt model file is present and readable."""
+    filepath = "densenet169_seed40_best.pt"
     if os.path.exists(filepath):
         return f"✅ Model file found at: {filepath}"
     else:
@@ -99,7 +63,7 @@ def test_file():
 
 @app.route('/predict', methods=['POST'])
 def predict():
-    """Perform prediction, save results (including uploaded image), and save to SQLite database."""
+    """Perform prediction using PyTorch, save results (including uploaded image), and save to SQLite database."""
     if 'file' not in request.files:
         return jsonify({'error': 'No file uploaded'}), 400
 
@@ -116,26 +80,24 @@ def predict():
 
         # ✅ Perform prediction
         img = Image.open(uploaded_file_path).convert('RGB')
-        img_array = preprocess_image(img)
+        input_tensor = transform(img).unsqueeze(0)
 
-        prediction = model.predict(img_array)[0]
+        with torch.no_grad():
+            prediction = model(input_tensor).numpy()[0]
+
+        # ✅ Interpret the output
         glaucoma_prob = 1 - prediction[0]
         normal_prob = prediction[0]
         result = 'Glaucoma' if glaucoma_prob > normal_prob else 'Normal'
         confidence = float(glaucoma_prob) if result == 'Glaucoma' else float(normal_prob)
 
-        # ✅ Grad-CAM
-        heatmap = make_gradcam(img_array, model, last_conv_layer_name='Conv_1_bn')
-        gradcam_filename = f"gradcam_{timestamp}.png"
-        save_gradcam_image(img, heatmap, filename=gradcam_filename)
-
-        # ✅ Save results to SQLite
+        # ✅ Save results to SQLite (no Grad-CAM generation for now)
         conn = sqlite3.connect(DB_PATH)
         cursor = conn.cursor()
         cursor.execute("""
             INSERT INTO results (image_filename, prediction, confidence, gradcam_filename, timestamp) 
             VALUES (?, ?, ?, ?, ?)
-        """, (uploaded_filename, result, confidence, gradcam_filename, datetime.now().isoformat()))
+        """, (uploaded_filename, result, confidence, '', datetime.now().isoformat()))
         conn.commit()
         conn.close()
 
@@ -144,13 +106,14 @@ def predict():
             'confidence': confidence,
             'normal_probability': float(normal_prob),
             'glaucoma_probability': float(glaucoma_prob),
-            'gradcam_image': gradcam_filename,
+            'gradcam_image': '',  # No Grad-CAM for PyTorch for now
             'image_filename': uploaded_filename
         })
 
     except Exception as e:
         return jsonify({'error': str(e)}), 500
 
+
 @app.route('/results', methods=['GET'])
 def results():
     """List all results from the SQLite database."""
@@ -175,7 +138,7 @@ def results():
 
 @app.route('/gradcam/<filename>')
 def get_gradcam(filename):
-    """Serve the Grad-CAM overlay image."""
+    """Serve the Grad-CAM overlay image (no-op if not used)."""
     filepath = os.path.join(OUTPUT_DIR, filename)
     if os.path.exists(filepath):
         return send_file(filepath, mimetype='image/png')
@@ -191,5 +154,6 @@ def get_image(filename):
     else:
         return jsonify({'error': 'File not found'}), 404
 
+
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860)