File size: 5,357 Bytes
b398c5e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | import tensorflow as tf
from tensorflow.keras.preprocessing import image
import numpy as np
import cv2
import os
from PIL import Image
# -----------------------------------------------------
# GLOBALS
# -----------------------------------------------------
cnn_model = None
gradcam_model = None
IMAGE_SIZE = (224, 224)
LAST_CONV_LAYER_NAME = 'conv2'
# -----------------------------------------------------
# LOAD MODEL & PREPARE GRAD-CAM
# -----------------------------------------------------
def load_cnn_model():
"""Loads the pre-trained CNN model and prepares Grad-CAM functional model."""
global cnn_model, gradcam_model
try:
if not os.path.exists('cnn_model.h5'):
raise FileNotFoundError("cnn_model.h5 is missing. Run create_dummy_cnn.py first.")
cnn_model = tf.keras.models.load_model('cnn_model.h5')
# CRITICAL FIX 1: Re-compile the model for reliable gradient calculation
cnn_model.compile(optimizer='adam', loss='binary_crossentropy', run_eagerly=False)
# Grad-CAM model outputs last conv layer + prediction
gradcam_model = tf.keras.models.Model(
inputs=cnn_model.input,
outputs=[cnn_model.get_layer(LAST_CONV_LAYER_NAME).output, cnn_model.output]
)
print(f"✅ CNN model loaded. Last Conv Layer: {LAST_CONV_LAYER_NAME}")
except Exception as e:
print(f"FATAL ERROR: Could not load or compile CNN model: {e}")
cnn_model = None
# Load immediately
load_cnn_model()
# -----------------------------------------------------
# IMAGE PROCESSING & PREDICTION
# -----------------------------------------------------
def get_img_array(img_path, size):
"""Utility function to load image and format it for the model."""
img = image.load_img(img_path, target_size=size)
array = image.img_to_array(img)
array = np.expand_dims(array, axis=0)
# CRITICAL FIX 2: Ensure array is float32 before normalization/prediction
array = tf.cast(array, dtype=tf.float32)
array /= 255.0
return array
def predict_xray_risk(img_path):
if cnn_model is None: return 0.0
try:
img_array = get_img_array(img_path, IMAGE_SIZE)
prediction = cnn_model.predict(img_array)[0]
return float(prediction[0])
except Exception as e:
print(f"CNN Prediction Error: {e}")
return 0.0
# -----------------------------------------------------
# GRAD-CAM FUNCTIONS (Ultimate Stability Fixes)
# -----------------------------------------------------
def make_gradcam_heatmap(img_path, pred_index=None):
"""Generates the Grad-CAM heatmap array with stability fixes."""
if gradcam_model is None: return np.zeros(IMAGE_SIZE[:2])
img_array = get_img_array(img_path, IMAGE_SIZE)
with tf.GradientTape() as tape:
last_conv_output, preds = gradcam_model(img_array)
if pred_index is None:
pred_index = tf.argmax(preds[0])
class_channel = preds[:, pred_index]
grads = tape.gradient(class_channel, last_conv_output)
pooled_grads = tf.reduce_mean(grads, axis=(0,1,2))
# Use stable multiplication for heatmap generation
last_conv_output_tensor = last_conv_output[0]
# Final stable multiplication
heatmap = last_conv_output_tensor * pooled_grads
heatmap = tf.reduce_sum(heatmap, axis=-1)
# Safe normalization
heatmap = tf.maximum(heatmap, 0)
max_val = tf.reduce_max(heatmap)
# CRITICAL FIX 3: Prevent division by zero if max_val is 0.0
if max_val == 0: max_val = tf.constant(1e-10, dtype=tf.float32)
heatmap /= max_val
return heatmap.numpy()
def save_gradcam_overlay(img_path, heatmap, save_filename, alpha=0.4):
"""Overlays the heatmap onto the original image and saves it to the static folder."""
static_dir = os.path.join(os.getcwd(), 'static')
os.makedirs(static_dir, exist_ok=True)
save_path = os.path.join(static_dir, save_filename)
# Core OpenCV Image Processing Logic (The stable part)
img = Image.open(img_path).convert('RGB')
heatmap = cv2.resize(heatmap, (img.width, img.height))
heatmap = np.uint8(255 * heatmap)
heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
# Use cv2.addWeighted for cleaner overlay math
superimposed_img = cv2.addWeighted(img_cv, 1-alpha, heatmap, alpha, 0)
cv2.imwrite(save_path, superimposed_img)
return save_filename
def generate_and_save_gradcam(original_xray_path):
"""Runs the full Grad-CAM pipeline and returns the FILENAME to the saved image."""
if gradcam_model is None: return None
base_name = os.path.basename(original_xray_path)
gradcam_filename = f"gradcam_{base_name}"
try:
heatmap = make_gradcam_heatmap(original_xray_path, pred_index=0)
saved_file = save_gradcam_overlay(original_xray_path, heatmap, gradcam_filename)
return saved_file
except Exception as e:
# Prints error to console if something failed during processing/saving
print(f"GRAD-CAM Generation Failed: {e}")
return None |