Update app.py

app.py

@@ -29,33 +29,49 @@ def extract_face(image):
     x, y = max(0, x), max(0, y)
     return image[y:y+h, x:x+w]
 
-def grad_cam(model, image, size, preprocess_fn):
-    img = cv2.resize(image, size)
-    input_tensor = preprocess_fn(img.astype(np.float32))[np.newaxis, ...]
-    
+def grad_cam(model, image, size, preprocess):
+    img_input = tf.expand_dims(image, axis=0)  # Shape: (1, H, W, 3)
+
+    last_conv_layer_name = None
+    # Auto-detect last conv layer
+    for layer in reversed(model.layers):
+        if isinstance(layer, tf.keras.layers.Conv2D):
+            last_conv_layer_name = layer.name
+            break
+
+    if not last_conv_layer_name:
+        raise ValueError("No Conv2D layer found in the model.")
+
     grad_model = tf.keras.models.Model(
-        [model.inputs], [model.get_layer(index=-1).output, model.output]
+        [model.inputs], [model.get_layer(last_conv_layer_name).output, model.output]
     )
 
     with tf.GradientTape() as tape:
-        conv_outputs, predictions = grad_model(input_tensor)
+        conv_outputs, predictions = grad_model(img_input)
         loss = predictions[:, 0]  # Assuming binary classification
 
-    grads = tape.gradient(loss, conv_outputs)[0]
-    conv_outputs = conv_outputs[0]
+    grads = tape.gradient(loss, conv_outputs)
+
+    # Ensure correct shape
+    if len(grads.shape) != 4:
+        raise ValueError(f"Expected 4D tensor for grads, got shape: {grads.shape}")
 
-    weights = tf.reduce_mean(grads, axis=(0, 1))
-    cam = np.zeros(conv_outputs.shape[:2], dtype=np.float32)
-    for i, w in enumerate(weights):
-        cam += w * conv_outputs[:, :, i]
+    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))  # (C,)
+    conv_outputs = conv_outputs[0]  # (H, W, C)
+    cam = tf.reduce_sum(conv_outputs * pooled_grads, axis=-1)
 
-    cam = np.maximum(cam, 0)
-    cam = cam / (cam.max() + 1e-8)
-    heatmap = cv2.resize(cam, size)
+    heatmap = np.maximum(cam, 0)
+    heatmap = heatmap / (tf.reduce_max(heatmap) + 1e-6)
+    heatmap = cv2.resize(heatmap.numpy(), size)
     heatmap = np.uint8(255 * heatmap)
     heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
-    superimposed_img = cv2.addWeighted(cv2.resize(img, size), 0.6, heatmap, 0.4, 0)
-    return superimposed_img
+
+    face_rgb = (image.numpy() * 255).astype(np.uint8)
+    face_rgb = cv2.resize(face_rgb, size)
+
+    overlaid = cv2.addWeighted(face_rgb, 0.6, heatmap, 0.4, 0)
+    return overlaid
+
 
 def predict(image):
     face = extract_face(image)