Update app.py

app.py

@@ -39,10 +39,11 @@ def grad_cam(model, img, size, preprocess_func):
         conv_outputs, predictions = grad_model(x_tensor)
         loss = predictions[:, 0]
     grads = tape.gradient(loss, conv_outputs)
-    cam = tf.reduce_mean(grads, axis=-1).numpy()[0]
+    cam = tf.reduce_mean(grads, axis=-1).numpy()
+    cam = np.squeeze(cam)
 
     cam = np.maximum(cam, 0)
-    cam /= cam.max() if cam.max() != 0 else 1
+    cam = cam / cam.max() if cam.max() != 0 else cam
     heatmap = cv2.resize(cam, size)
     heatmap = np.uint8(255 * heatmap)
     heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
@@ -54,6 +55,7 @@ def grad_cam(model, img, size, preprocess_func):
     return Image.fromarray(cv2.cvtColor(superimposed, cv2.COLOR_BGR2RGB))
 
 
+
 # Preprocessing helper
 def preprocess(img, size, func):
     img = img.resize(size)