Update app.py

app.py

@@ -238,6 +238,10 @@ def FSFM3C_image_detection(image):
         }
         target_category = TargetCategory(category_names_to_index[max_prob_class])
 
+        cam = GradCAM(model=model,
+                      target_layers=[model.blocks[-1].norm1],
+                      reshape_transform=reshape_transform
+                      )
         grayscale_cam = cam(input_tensor=input_tensor, targets=[target_category])
         grayscale_cam = 1 - grayscale_cam[0, :]
         img = np.array(img)
@@ -250,7 +254,10 @@ def FSFM3C_image_detection(image):
         # Add text overlay to the heatmap
         # text = f"Detected: {max_prob_class}"
         # cv2.putText(visualization, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
-        output_path = "./CAM_images/output_heatmap.png"
+        cam_path = os.path.join(CAM_SAVE_PATH, str(len(os.listdir(CAM_SAVE_PATH))))
+        os.makedirs(cam_path, exist_ok=True)
+        os.makedirs(os.path.join(cam_path, '0'), exist_ok=True)
+        output_path = os.path.join("output_heatmap.png")
         cv2.imwrite(output_path, visualization)
         return image_results, output_path, probabilities[max_prob_index]
 
@@ -329,8 +336,10 @@ def FSFM3C_video_detection(video, num_frames):
 # Paths and Constants
 P = os.path.abspath(__file__)
 FRAME_SAVE_PATH = os.path.join(os.path.dirname(P), 'frame')
+CAM_SAVE_PATH = os.path.join(os.path.dirname(P), 'cam')
 CKPT_SAVE_PATH = os.path.join(os.path.dirname(P), 'checkpoints')
 os.makedirs(FRAME_SAVE_PATH, exist_ok=True)
+os.makedirs(CAM_SAVE_PATH, exist_ok=True)
 os.makedirs(CKPT_SAVE_PATH, exist_ok=True)
 CKPT_NAME = [
     '✨Unified-detector_v1_Fine-tuned_on_4_classes',