Update app.py

app.py

@@ -77,18 +77,48 @@ def process_image(image_path):
 
     # Check for anomaly
     if is_anomaly(clf, feature_extractor, input_image):
-        return "Anomaly detected. Image will not be classified.", None, None
+        return "Anomaly detected. Image will not be classified.", None
 
     # Classify image
     predicted_class, probability = classify_image(classification_model, input_image)
     result = f"The predicted class is: {predicted_class} with a probability of {probability:.2f}%"
 
-    # Further processing for heatmap or additional features can be added here
-
-    return result, None, None  # Returning placeholders for additional outputs if needed
+    # Generate heatmap
+    heatmap = generate_heatmap(classification_model, input_image)
+    heatmap_image = Image.fromarray(np.uint8(plt.cm.hot(heatmap) * 255))
+
+    return result, heatmap_image
+
+# Function to generate heatmap
+def generate_heatmap(model, image):
+    activation = []
+    def hook_fn(module, input, output):
+        activation.append(output)
+    for module in model.named_modules():
+        if isinstance(module[1], torch.nn.ReLU):
+            module[1].register_forward_hook(hook_fn)
+
+    # Forward pass
+    output = model(image)
+    prediction = output.argmax(1)
+
+    # Backpropagation to compute gradients
+    model.zero_grad()
+    one_hot_output = torch.FloatTensor(1, output.size()[-1]).zero_().to(device)
+    one_hot_output[0][prediction] = 1
+    output.backward(gradient=one_hot_output)
+
+    # Compute the heatmap
+    if len(activation) > 0:
+        gradients = activation[0][0].detach().cpu().numpy()
+        heatmap = gradients.max(axis=0)
+        threshold = 0.5  # Adjust this threshold value as needed
+        heatmap[heatmap < threshold] = 0
+        return heatmap
+    else:
+        return np.zeros((224, 224))  # Return an empty heatmap if no activation is recorded
 
 # Gradio interface
-    
 iface = gr.Interface(
     fn=process_image,
     inputs=File(type="filepath"),
@@ -97,7 +127,7 @@ iface = gr.Interface(
     description="Upload an image to classify it as normal or abnormal.",
     article="Above is a sample image to test the results of the model. Click it to see the results.",
     examples=[
-        ["Gastric_Images/Ladybug.png"],
+        ["/Gastric_Images/Ladybug.png"],
     ],
     allow_flagging="never",
 )