Update app.py

app.py

@@ -1,8 +1,7 @@
 import os
-# --- FIX: Disable GPU to prevent CUDA initialization errors ---
 os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
 
-from flask import Flask, render_template, request, jsonify, session, redirect, url_for, send_from_directory
+from flask import Flask, render_template, request, jsonify, redirect, url_for, send_from_directory
 from flask_pymongo import PyMongo
 from flask_bcrypt import Bcrypt
 import tensorflow as tf
@@ -17,30 +16,20 @@ import uuid
 import secrets
 import logging
 
-# -------------------- Setup & Config --------------------
-
-# Load environment variables
 load_dotenv()
 
 app = Flask(__name__)
 
-# Configurations
 app.config["MONGO_URI"] = os.getenv("MONGODB_URI") or os.getenv("MONGO_URI")
-# Keep your format, just ensure it's set once per process
 app.config['SECRET_KEY'] = os.getenv("SECRET_KEY") or secrets.token_hex(16)
-
-# Slightly safer cookie defaults without changing your session usage
 app.config.setdefault("SESSION_COOKIE_HTTPONLY", True)
 app.config.setdefault("SESSION_COOKIE_SAMESITE", "Lax")
 
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
 
-# Basic logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("app")
 
-# Initialize extensions
-# If MONGO_URI missing, still construct PyMongo but avoid immediate use crashes
 try:
     if app.config["MONGO_URI"]:
         mongo = PyMongo(app, tlsCAFile=certifi.where())
@@ -49,24 +38,20 @@ try:
         mongo = PyMongo(app, tlsCAFile=certifi.where())
 except Exception as e:
     logger.error(f"Mongo initialization error: {e}")
-    # Keep an object to avoid NameError later
     mongo = None
 
 bcrypt = Bcrypt(app)
 
-# Configure Gemini
-gemini_model = "gemini-2.0-flash"
+gemini_model = None
 if GEMINI_API_KEY:
     try:
         genai.configure(api_key=GEMINI_API_KEY)
-        # Keep your model name
         gemini_model = genai.GenerativeModel('gemini-2.0-flash')
     except Exception as e:
         logger.error(f"Gemini initialization error: {e}")
 else:
     logger.warning("GEMINI_API_KEY/GOOGLE_API_KEY not set. /chat will return a friendly error.")
 
-# --- Model Configuration ---
 MODEL_CONFIG = {
     "Pneumonia": {
         "path": "model/best_pneumonia_model.h5",
@@ -102,11 +87,9 @@ MODEL_CONFIG = {
     }
 }
 
-# --- Model Loading ---
 models = {}
 
 def load_all_models():
-    """Loads all models from the 'model' directory based on MODEL_CONFIG."""
     for name, config in MODEL_CONFIG.items():
         try:
             model_path = config["path"]
@@ -118,15 +101,11 @@ def load_all_models():
         except Exception as e:
             logger.error(f"Error loading model {name}: {e}")
 
-# Load models on application startup
 load_all_models()
 
-# --- Image Preprocessing ---
 def preprocess_image(img_path, target_size=(224, 224)):
-    """Preprocesses the image for model prediction."""
     img = image.load_img(img_path, target_size=target_size)
     img_array = image.img_to_array(img)
-    # Handle grayscale or alpha automatically by broadcasting if needed
     if img_array.ndim == 2:
         img_array = np.stack([img_array]*3, axis=-1)
     elif img_array.shape[-1] == 4:
@@ -135,102 +114,68 @@ def preprocess_image(img_path, target_size=(224, 224)):
     img_array = img_array.astype("float32") / 255.0
     return img_array
 
-# --- Grad-CAM Utilities ---
 def _safe_get_layer(model, layer_name):
-    """Return layer if exists; else None."""
     try:
         return model.get_layer(layer_name)
     except Exception:
         return None
 
 def find_last_conv_layer(model):
-    """Finds the name of the last convolutional layer in a model."""
-    logger.info("--- DEBUG: Searching for last convolutional layer ---")
     for layer in reversed(model.layers):
         if isinstance(layer, (tf.keras.layers.Conv2D, tf.keras.layers.DepthwiseConv2D)):
-            # 4D output: (batch, h, w, channels)
             try:
                 out_shape = layer.output_shape
             except Exception:
                 out_shape = None
             if out_shape and len(out_shape) == 4:
-                logger.info(f"Found candidate last conv layer: {layer.name}")
                 return layer.name
     raise ValueError("Could not automatically find a convolutional layer in the model.")
 
 def get_gradcam_heatmap(model, img_array, last_conv_layer_name, pred_index=None):
-    """Generates a Grad-CAM heatmap."""
-    # If configured layer isn't present, auto-detect
     if not _safe_get_layer(model, last_conv_layer_name):
         last_conv_layer_name = find_last_conv_layer(model)
-
     conv_layer = model.get_layer(last_conv_layer_name)
-    grad_model = tf.keras.models.Model(
-        [model.inputs], [conv_layer.output, model.output]
-    )
-
+    grad_model = tf.keras.models.Model([model.inputs], [conv_layer.output, model.output])
     with tf.GradientTape() as tape:
         conv_outputs, preds = grad_model(img_array, training=False)
-
         if isinstance(preds, (list, tuple)):
             preds = preds[0]
-
-        # Ensure preds is a tensor
         preds = tf.convert_to_tensor(preds)
-
-        # If model is binary with single logit/sigmoid output
         if preds.shape.rank is not None and preds.shape[-1] == 1:
             class_channel = preds[:, 0]
         else:
             if pred_index is None:
                 pred_index = tf.argmax(preds[0])
             class_channel = preds[:, pred_index]
-
     grads = tape.gradient(class_channel, conv_outputs)
     if grads is None:
-        # Fallback: no gradient (e.g., custom layers). Return uniform zeros heatmap.
         heatmap = tf.zeros(conv_outputs.shape[1:3], dtype=tf.float32)
         return heatmap.numpy()
-
     pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
     conv_outputs = conv_outputs[0]
     heatmap = tf.tensordot(conv_outputs, pooled_grads, axes=(2, 0))
-
     heatmap = tf.maximum(heatmap, 0)
     denom = tf.math.reduce_max(heatmap)
     heatmap = heatmap / (denom + 1e-8)
     return heatmap.numpy()
 
 def save_gradcam_image(img_path, heatmap, output_path, threshold=0.6, alpha=0.4):
-    """
-    Saves the Grad-CAM image by highlighting only the most important areas
-    with light red spots.
-    """
     img = cv2.imread(img_path)
     if img is None:
         raise ValueError("Failed to read image with OpenCV.")
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # Work with RGB
-
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]))
-
-    # Create a mask where the heatmap is above the threshold
     mask = heatmap > threshold
-
-    # Create a red overlay
     overlay = np.zeros_like(img, dtype=np.uint8)
-    overlay[mask] = [255, 0, 0]  # Red color for highlighted spots
-
-    # Blend the original image with the red overlay using the mask
+    overlay[mask] = [255, 0, 0]
     superimposed_img = cv2.addWeighted(overlay, alpha, img, 1 - alpha, 0)
-
-    # Areas outside the mask should be the original image
     superimposed_img[~mask] = img[~mask]
-
     superimposed_img = cv2.cvtColor(superimposed_img, cv2.COLOR_RGB2BGR)
     cv2.imwrite(output_path, superimposed_img)
     return output_path
 
-# -------------------- Routes --------------------
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+TEST_IMAGES_DIR = os.path.join(BASE_DIR, 'testimages')
 
 @app.route("/")
 def home():
@@ -240,44 +185,48 @@ def home():
 def serve_tmp_file(filename):
     return send_from_directory('/tmp', filename)
 
+@app.route('/testimages/<path:filename>')
+def serve_test_image(filename):
+    return send_from_directory(TEST_IMAGES_DIR, filename)
+
+@app.route('/example_images')
+def example_images():
+    try:
+        files = []
+        if os.path.isdir(TEST_IMAGES_DIR):
+            for f in os.listdir(TEST_IMAGES_DIR):
+                lf = f.lower()
+                if lf.endswith(('.png', '.jpg', '.jpeg')):
+                    files.append(url_for('serve_test_image', filename=f))
+        return jsonify({"images": files})
+    except Exception as e:
+        logger.error(f"example_images error: {e}")
+        return jsonify({"images": []})
+
 @app.route('/login', methods=['GET', 'POST'])
 def login():
-    # Authentication removed; redirect to main app
     return redirect(url_for('index'))
 
 @app.route('/signup', methods=['GET', 'POST'])
 def signup():
-    # Authentication removed; redirect to main app
     return redirect(url_for('index'))
 
 @app.route('/index')
 def index():
-    # Publicly accessible index
     return render_template('index.html')
 
 @app.route('/logout')
 def logout():
-    # Authentication removed; redirect to main app
     return redirect(url_for('index'))
 
 def _postprocess_binary_prediction(raw):
-    """
-    Normalize binary outputs across shapes:
-    - (1,) or (N,)      : sigmoid probabilities
-    - (1,1) or (N,1)    : sigmoid probabilities
-    - logits also supported (auto-sigmoid)
-    Returns probability in [0,1].
-    """
     arr = np.array(raw, dtype=np.float32)
     arr = np.squeeze(arr)
-    # If scalar, keep it
     if arr.ndim == 0:
         prob = float(arr)
-        # Heuristic: if obviously a logit (|x|>1 and not in [0,1]), apply sigmoid
         if prob < 0.0 or prob > 1.0:
             prob = float(1.0 / (1.0 + np.exp(-prob)))
         return min(max(prob, 0.0), 1.0)
-    # If 1D vector, take first
     prob = float(arr[0])
     if prob < 0.0 or prob > 1.0:
         prob = float(1.0 / (1.0 + np.exp(-prob)))
@@ -287,33 +236,23 @@ def _postprocess_binary_prediction(raw):
 def predict():
     if "file" not in request.files:
         return jsonify({"error": "No file part"}), 400
-
     file = request.files["file"]
     model_name = request.form.get("model")
-
     if not file or file.filename == "":
         return jsonify({"error": "No selected file"}), 400
-
     if model_name not in models:
         return jsonify({"error": "Invalid model selected"}), 400
-
     try:
         filename = f"{uuid.uuid4()}_{file.filename}"
         filepath = os.path.join("/tmp", filename)
         file.save(filepath)
-
         model_config = MODEL_CONFIG[model_name]
         model = models[model_name]
         labels = model_config["labels"]
         input_size = model_config.get("input_size", (224, 224))
-
         img_array = preprocess_image(filepath, target_size=input_size)
         prediction = model.predict(img_array, verbose=0)
-
-        # Ensure numpy array
         prediction = np.array(prediction)
-
-        # Binary case (2 labels) with single neuron output (logit or sigmoid)
         if len(labels) == 2 and prediction.ndim >= 1 and prediction.shape[-1] in (1,) and prediction.size >= 1:
             prob_pos = _postprocess_binary_prediction(prediction)
             if prob_pos >= 0.5:
@@ -325,12 +264,10 @@ def predict():
                 predicted_label = labels[0]
                 confidence = 1.0 - prob_pos
         else:
-            # Multi-class: softmax or logits
             if prediction.ndim == 2:
                 vec = prediction[0]
             else:
                 vec = prediction.reshape(-1)
-            # If appears to be logits, apply softmax for confidence; otherwise trust as probs
             if np.any(vec < 0) or np.any(vec > 1) or not np.isclose(np.sum(vec), 1.0, atol=1e-3):
                 exps = np.exp(vec - np.max(vec))
                 probs = exps / (np.sum(exps) + 1e-8)
@@ -339,25 +276,16 @@ def predict():
             predicted_index = int(np.argmax(probs))
             predicted_label = labels[predicted_index]
             confidence = float(np.max(probs))
-
         gradcam_url = None
         try:
-            logger.info(f"--- Generating Grad-CAM for model: {model_name} ---")
             last_conv_layer_name = MODEL_CONFIG[model_name].get('last_conv_layer') or ""
             heatmap = get_gradcam_heatmap(model, img_array, last_conv_layer_name, pred_index=predicted_index)
-
             gradcam_filename = f"gradcam_{filename}"
             gradcam_filepath = os.path.join("/tmp", gradcam_filename)
             save_gradcam_image(filepath, heatmap, gradcam_filepath)
             gradcam_url = url_for('serve_tmp_file', filename=gradcam_filename)
-            logger.info("--- Successfully generated Grad-CAM image ---")
         except Exception as e:
-            logger.error(f"--- Grad-CAM Generation FAILED for model: {model_name} --- Error: {e}")
-            try:
-                model.summary(print_fn=lambda x: logger.info(x))
-            except Exception:
-                pass
-
+            logger.error(f"Grad-CAM error: {e}")
         return jsonify({
             "original_image": url_for('serve_tmp_file', filename=filename),
             "gradcam_image": gradcam_url,
@@ -374,8 +302,6 @@ def chat():
     data = request.get_json(silent=True) or {}
     user_message = data.get("message", "")
     prediction_context = data.get("context") or {}
-
-    # Guard against missing keys
     model_used = prediction_context.get('model_used', 'Unknown Model')
     pred_label = prediction_context.get('prediction', 'Unknown')
     conf = prediction_context.get('confidence', 0.0)
@@ -383,7 +309,6 @@ def chat():
         conf_pct = float(conf) * 100.0
     except Exception:
         conf_pct = 0.0
-
     prompt = f"""
 You are a helpful medical assistant chatbot.
 A medical image was analyzed with the following results:
@@ -394,20 +319,16 @@ The user's question is: "{user_message}"
 Based on this context, provide a helpful and informative response.
 Do not provide a diagnosis. Advise the user to consult a medical professional.
 """
-
     try:
         if gemini_model is None:
             return jsonify({"error": "Gemini API not configured. Set GEMINI_API_KEY in environment."}), 500
         response = gemini_model.generate_content(prompt)
-        # Some SDKs return .text; guard if attribute missing
         text = getattr(response, "text", None)
         if not text:
-            # Try to stringify safely
             text = str(response)
         return jsonify({"response": text})
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
 if __name__ == "__main__":
-    # Keep your debug flag as-is
     app.run(debug=True)