Update app.py

app.py

@@ -23,6 +23,10 @@ labels = [
 # Function to apply Grad-CAM visualization
 def generate_grad_cam(image, target_layer):
     try:
+        # Convert image to RGB if necessary
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
         # Preprocess the image
         inputs = image_processor(images=image, return_tensors="pt")
         
@@ -45,14 +49,19 @@ def generate_grad_cam(image, target_layer):
         cam = torch.mean(pooled_gradients * inputs["pixel_values"], dim=1).squeeze()
         cam = torch.clamp(cam, min=0).numpy()  # Ensure non-negative values
         
-        return cam, predicted_class.item()
+        return cam, predicted_class.item(), None  # No error
     except Exception as e:
-        print(f"Error generating Grad-CAM: {e}")
-        return None
+        error_message = f"Error generating Grad-CAM: {e}"
+        print(error_message)
+        return None, None, error_message
 
 # Function to predict classes and visualize Grad-CAM
 def predict_and_explain(image):
     try:
+        # Convert image to RGB if necessary
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
         # Preprocess the image
         inputs = image_processor(images=image, return_tensors="pt")
         
@@ -62,33 +71,54 @@ def predict_and_explain(image):
         logits = outputs.logits
         
         predicted_class = logits.argmax(-1).item()
-        cam_map, _ = generate_grad_cam(image, "pooler_output")
+        cam_map, _, grad_cam_error = generate_grad_cam(image, "pooler_output")
+        
+        # Check for Grad-CAM errors
+        if grad_cam_error is not None:
+            return {
+                "predicted class": "Error during Grad-CAM generation",
+                "Grad-CAM map": None,
+                "error log": grad_cam_error
+            }
         
         # Convert cam_map to a visualizable format (heatmap)
-        plt.imshow(cam_map, cmap='jet', alpha=0.5)
-        plt.axis('off')
-        plt.title(f"Grad-CAM for {labels[predicted_class]}")
-        plt.colorbar()
-        plt.savefig("grad_cam_output.png")
-        plt.close()
+        if cam_map is not None:
+            plt.imshow(cam_map, cmap='jet', alpha=0.5)
+            plt.axis('off')
+            plt.title(f"Grad-CAM for {labels[predicted_class]}")
+            plt.colorbar()
+            plt.savefig("grad_cam_output.png")
+            plt.close()
         
-        # Load the saved image to return it
-        grad_cam_image = Image.open("grad_cam_output.png")
+            # Load the saved image to return it
+            grad_cam_image = Image.open("grad_cam_output.png")
+        else:
+            grad_cam_image = None
         
         return {
             "predicted class": labels[predicted_class],
             "Grad-CAM map": grad_cam_image,
+            "error log": "No errors"
         }
     except Exception as e:
-        print(f"Error predicting and explaining: {e}")
-        return None
+        error_message = f"Error predicting and explaining: {e}"
+        print(error_message)
+        return {
+            "predicted class": "Error during prediction",
+            "Grad-CAM map": None,
+            "error log": error_message
+        }
 
-# Create a Gradio interface
+# Create a Gradio interface with an error log output
 iface = gr.Interface(
     fn=predict_and_explain, 
     inputs="image", 
-    outputs=["text", "image"], 
-    title="Chest X-ray Classification"
+    outputs=[
+        gr.outputs.Textbox(label="Predicted Class"),
+        gr.outputs.Image(label="Grad-CAM Map"),
+        gr.outputs.Textbox(label="Error Log")  # Error log for debugging
+    ], 
+    title="Chest X-ray Classification with Debugging Logs"
 )
 
 if __name__ == "__main__":