Update app.py

app.py

@@ -24,7 +24,7 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 def setup_model():
     cfg = get_cfg()
     cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-    #cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
+    cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")
     cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")
 
     cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # Set a threshold for this model
@@ -99,15 +99,6 @@ def overlay_heatmap_on_image(heatmap, image, alpha=0.4, colormap=cv2.COLORMAP_JE
     overlayed_img = cv2.addWeighted(image, alpha, heatmap, 1 - alpha, 0)
     return overlayed_img
 
-# # Set up the model
-# def setup_model():
-#     cfg = get_cfg()
-#     cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
-#     cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # set threshold for this model
-#     cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
-#     cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"  # Use GPU if available
-#     predictor = DefaultPredictor(cfg)
-#     return predictor
 
 # Function to segment image
 def segment_image(image):
@@ -213,28 +204,28 @@ def segment_image(image):
     modernity_scores = calculate_modernity_scores(modernity_output, year_categories).item()
 
 
-    target_layer = modernity_model.layer4[-1]
+    # target_layer = modernity_model.layer4[-1]
 
-    modernity_cam = GradCAM(modernity_model, target_layer)
-    modernity_heatmap = modernity_cam(input_tensor, class_idx=torch.argmax(modernity_output).item())
+    # modernity_cam = GradCAM(modernity_model, target_layer)
+    # modernity_heatmap = modernity_cam(input_tensor, class_idx=torch.argmax(modernity_output).item())
 
-    target_layer = typicality_model.layer4[-1]
-    typicality_cam = GradCAM(typicality_model, target_layer)
-    typicality_heatmap = typicality_cam(input_tensor, class_idx=torch.argmax(typicality_output).item())
+    # target_layer = typicality_model.layer4[-1]
+    # typicality_cam = GradCAM(typicality_model, target_layer)
+    # typicality_heatmap = typicality_cam(input_tensor, class_idx=torch.argmax(typicality_output).item())
 
-    # Convert the input image to a format suitable for overlaying
-    img_np = np.array(cropped_car_pil)
-    img_np = cv2.resize(img_np, (224, 224))
+    # # Convert the input image to a format suitable for overlaying
+    # img_np = np.array(cropped_car_pil)
+    # img_np = cv2.resize(img_np, (224, 224))
 
-    # Overlay the heatmap on the image
-    overlayed_img_modernity = overlay_heatmap_on_image(modernity_heatmap, img_np)
-    overlayed_img_typicality = overlay_heatmap_on_image(typicality_heatmap, img_np)
+    # # Overlay the heatmap on the image
+    # overlayed_img_modernity = overlay_heatmap_on_image(modernity_heatmap, img_np)
+    # overlayed_img_typicality = overlay_heatmap_on_image(typicality_heatmap, img_np)
 
-    # Convert overlayed images back to PIL for saving
-    overlayed_img_modernity_pil = Image.fromarray(cv2.cvtColor(overlayed_img_modernity, cv2.COLOR_BGR2RGB))
-    overlayed_img_typicality_pil = Image.fromarray(cv2.cvtColor(overlayed_img_typicality, cv2.COLOR_BGR2RGB))
+    # # Convert overlayed images back to PIL for saving
+    # overlayed_img_modernity_pil = Image.fromarray(cv2.cvtColor(overlayed_img_modernity, cv2.COLOR_BGR2RGB))
+    # overlayed_img_typicality_pil = Image.fromarray(cv2.cvtColor(overlayed_img_typicality, cv2.COLOR_BGR2RGB))
 
-    return "Automobiles detected in the image", cropped_car_pil, modernity_scores, typicality_scores , most_similar_group, overlayed_img_modernity_pil,  overlayed_img_typicality_pil
+    return "Automobiles detected in the image", cropped_car_pil, modernity_scores, typicality_scores , most_similar_group#, overlayed_img_modernity_pil,  overlayed_img_typicality_pil
 
 
 # Create Gradio interface