Update app.py

app.py

@@ -40,13 +40,8 @@ np.random.seed(SEED)
 
 
 
-# Get the current working directory
-current_dir = os.path.dirname(__file__)
-# Construct the path to the CSV file
-file_path = os.path.join(current_dir, 'morphs.csv')
 
-
-df = pd.read_csv(file_path)
+df = pd.read_csv('morphs.csv')
 
 morphs = {}
 for _, row in df.iterrows():
@@ -100,6 +95,15 @@ 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):
@@ -141,7 +145,6 @@ def segment_image(image):
 
     # Preprocess the cropped car image for the model
     transform = transforms.Compose([
-        transforms.Resize(256),
         transforms.Resize((224, 224)),
         transforms.ToTensor(),
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
@@ -205,26 +208,26 @@ 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
 
@@ -239,8 +242,8 @@ iface = gr.Interface(
         gr.Textbox(label="Modernity Score"),
         gr.Textbox(label="Typicality Score"),
         gr.Textbox(label="most_similar_group"),
-        gr.Image(label="Grad-CAM for Year"),
-        gr.Image(label="Grad-CAM for Type")
+        gr.Image(label="Grad-CAM for Type"),
+        gr.Image(label="Grad-CAM for Year")
     ],
     title="Automobile Detection and Scoring using Mask R-CNN",
     description="Upload an image, and the system will detect and segment automobiles, crop the largest car, and predict its modernity and typicality scores. Grad-CAM heatmaps will also be generated."