Update

app.py

@@ -1,35 +1,41 @@
 import gradio as gr
 import predict as predict
-from googletrans import Translator, constants
-from pprint import pprint
-
-translator = Translator()
-
+import extract as extract
 
 def upload_file(files):
     file_paths = [file.name for file in files]
     return file_paths
 
 
-def process_file(webcam_filepath, upload_filepath):
+def process_file(webcam_filepath, upload_filepath, ):
     result = []
     if webcam_filepath == None:
         sorted_classes = predict.predict_image(upload_filepath)
         for class_label, class_prob in sorted_classes:
             class_prob = class_prob.item().__round__(2)
             result.append(f"{class_label}: {class_prob}%")
+        cam = extract.extract_gradcam(upload_filepath, save_path="gradcam.jpg")
+        result = result[:3]
+        result.append("gradcam.jpg")
         return result
     elif upload_filepath == None:
         sorted_classes = predict.predict_image(webcam_filepath)
         for class_label, class_prob in sorted_classes:
             class_prob = class_prob.item().__round__(2)
             result.append(f"{class_label}: {class_prob}%")
+        cam = extract.extract_gradcam(webcam_filepath, save_path="gradcam.jpg")
+        result = result[:3]
+        result.append("gradcam.jpg")
         return result
     else:
         sorted_classes = predict.predict_image(upload_filepath)
         for class_label, class_prob in sorted_classes:
             class_prob = class_prob.item().__round__(2)
             result.append(f"{class_label}: {class_prob}%")
+        cam = extract.extract_gradcam(upload_filepath, save_path="gradcam.jpg")
+        # Only keep the first 3 results
+        result = result[:3]
+        result.append("gradcam.jpg")
         return result
 
 
@@ -46,7 +52,12 @@ demo = gr.Interface(
         gr.outputs.Textbox(label="Probability 1"),
         gr.outputs.Textbox(label="Probability 2"),
         gr.outputs.Textbox(label="Probability 3"),
+        # GradCAM
+        gr.outputs.Image(label="GradCAM++", type="filepath"),
+        
+        
     ],
+    
 )
 
 demo.launch()

extract.py

@@ -8,10 +8,9 @@ from configs import *
 import os, random
 
 # Load your model (replace with your model class)
-model = MODEL  # Replace with your model
+model = MODEL.to(DEVICE)
 model.load_state_dict(torch.load(MODEL_SAVE_PATH))
 model.eval()
-model = model.to(DEVICE)
 
 # Find the target layer (modify this based on your model architecture)
 target_layer = None
@@ -22,33 +21,63 @@ for child in model.features[-1]:
 if target_layer is None:
     raise ValueError("Invalid layer name: {}".format(target_layer))
 
+def extract_gradcam(image_path=None, save_path=None):
+    if image_path is None:
+        for disease in CLASSES:
+            print("Processing", disease)
+            image_path = random.choice(os.listdir("data/test/Task 1/" + disease))
+            image_path = "data/test/Task 1/" + disease + "/" + image_path
+            rgb_img = cv2.imread(image_path, 1)
+            rgb_img = np.float32(rgb_img) / 255
+            input_tensor = preprocess_image(rgb_img, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+            input_tensor = input_tensor.to(DEVICE)
 
-for disease in CLASSES:
-    print("Processing", disease)
-    image_path = random.choice(os.listdir("data/test/Task 1/" + disease))
-    image_path = "data/test/Task 1/" + disease + "/" + image_path
-    rgb_img = cv2.imread(image_path, 1)
-    rgb_img = np.float32(rgb_img) / 255
-    input_tensor = preprocess_image(rgb_img, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-    input_tensor = input_tensor.to(DEVICE)
+            # Create a GradCAMPlusPlus object
+            cam = GradCAMPlusPlus(model=model, target_layers=[target_layer], use_cuda=True)
 
-    # Create a GradCAMPlusPlus object
-    cam = GradCAMPlusPlus(model=model, target_layers=[target_layer], use_cuda=True)
+            # Generate the GradCAM heatmap
+            grayscale_cam = cam(input_tensor=input_tensor)[0]
 
-    # Generate the GradCAM heatmap
-    grayscale_cam = cam(input_tensor=input_tensor)[0]
+            # Apply a colormap to the grayscale heatmap
+            heatmap_colored = cv2.applyColorMap(np.uint8(255 * grayscale_cam), cv2.COLORMAP_JET)
 
-    # Apply a colormap to the grayscale heatmap
-    heatmap_colored = cv2.applyColorMap(np.uint8(255 * grayscale_cam), cv2.COLORMAP_JET)
+            # Ensure heatmap_colored has the same dtype as rgb_img
+            heatmap_colored = heatmap_colored.astype(np.float32) / 255
 
-    # Ensure heatmap_colored has the same dtype as rgb_img
-    heatmap_colored = heatmap_colored.astype(np.float32) / 255
+            # Adjust the alpha value to control transparency
+            alpha = 0.3  # You can change this value to make the original image more or less transparent
 
-    # Adjust the alpha value to control transparency
-    alpha = 0.3  # You can change this value to make the original image more or less transparent
+            # Overlay the colored heatmap on the original image
+            final_output = cv2.addWeighted(rgb_img, 0.3, heatmap_colored, 0.7, 0)
 
-    # Overlay the colored heatmap on the original image
-    final_output = cv2.addWeighted(rgb_img, 0.3, heatmap_colored, 0.7, 0)
+            # Save the final output
+            cv2.imwrite(f'docs/efficientnet/gradcam/{disease}.jpg', (final_output * 255).astype(np.uint8))
+    else:
+            rgb_img = cv2.imread(image_path, 1)
+            rgb_img = np.float32(rgb_img) / 255
+            input_tensor = preprocess_image(rgb_img, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+            input_tensor = input_tensor.to(DEVICE)
 
-    # Save the final output
-    cv2.imwrite(f'docs/efficientnet/gradcam/{disease}.jpg', (final_output * 255).astype(np.uint8))
+            # Create a GradCAMPlusPlus object
+            cam = GradCAMPlusPlus(model=model, target_layers=[target_layer], use_cuda=True)
+
+            # Generate the GradCAM heatmap
+            grayscale_cam = cam(input_tensor=input_tensor)[0]
+
+            # Apply a colormap to the grayscale heatmap
+            heatmap_colored = cv2.applyColorMap(np.uint8(255 * grayscale_cam), cv2.COLORMAP_JET)
+
+            # Ensure heatmap_colored has the same dtype as rgb_img
+            heatmap_colored = heatmap_colored.astype(np.float32) / 255
+
+            # Adjust the alpha value to control transparency
+            alpha = 0.3  # You can change this value to make the original image more or less transparent
+
+            # Overlay the colored heatmap on the original image
+            final_output = cv2.addWeighted(rgb_img, 0.3, heatmap_colored, 0.7, 0)
+
+            # Save the final output
+            cv2.imwrite(save_path, (final_output * 255).astype(np.uint8))
+            
+            return save_path            
+