[+] plots evaluation and ranks

MathieuNlp · Oct 4, 2023 · c454007 · c454007
1 parent 941d3de
commit c454007
Show file tree

Hide file tree

Showing 44 changed files with 115 additions and 163 deletions.
diff --git a/annotations.json b/annotations.json
@@ -5,11 +5,14 @@
 "ring2.jpg": {"mask_path":"./dataset/train/masks/ring2.jpg", "bbox": [569, 239, 1286, 944]},
 "ring3.jpg": {"mask_path":"./dataset/train/masks/ring3.jpg", "bbox": [135, 68, 466, 415]},
 "ring4.jpg": {"mask_path":"./dataset/train/masks/ring4.jpg", "bbox": [133, 234, 1085, 718]},
-"ring5.jpg": {"mask_path":"./dataset/train/masks/ring5.jpg", "bbox": [163, 197, 477, 597]}
+"ring5.jpg": {"mask_path":"./dataset/train/masks/ring5.jpg", "bbox": [163, 197, 477, 597]},
+"ring6.jpg": {"mask_path":"./dataset/train/masks/ring6.jpg", "bbox": [12, 30, 438, 288]},
+"ring7.jpg": {"mask_path":"./dataset/train/masks/ring7.jpg", "bbox": [4, 11, 789, 438]},
+"ring8.jpg": {"mask_path":"./dataset/train/masks/ring8.jpg", "bbox": [185, 83, 621, 639]}
 },
 "test": 
-{"ring_test_1.jpg" : {"bbox": [74, 351, 936, 649]},
-"ring_test_2.jpg" : {"bbox": [138, 196, 508, 610]}
+{"ring_test_1.jpg" : {"mask_path":"./dataset/test/masks/ring_test_1.jpg", "bbox": [74, 351, 936, 649]},
+"ring_test_2.jpg" : {"mask_path":"./dataset/test/masks/ring_test_2.jpg", "bbox": [138, 196, 508, 610]}
 }
 
 }
diff --git a/eval_inference.py b/eval_inference.py
diff --git a/inference_baseline.py b/inference_baseline.py
diff --git a/inference_eval.py b/inference_eval.py
@@ -0,0 +1,99 @@
+import torch
+import monai
+from tqdm import tqdm
+from statistics import mean
+from torch.utils.data import Dataset, DataLoader
+from torchvision import datasets, transforms
+from torch.optim import Adam
+from torch.nn.functional import threshold, normalize
+from torchvision.utils import save_image
+import src.utils as utils
+from src.dataloader import DatasetSegmentation, collate_fn
+from src.processor import Samprocessor
+from src.segment_anything import build_sam_vit_b, SamPredictor
+from src.lora import LoRA_sam
+import matplotlib.pyplot as plt
+import yaml
+import torch.nn.functional as F
+import monai
+import numpy as np
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+seg_loss = monai.losses.DiceCELoss(sigmoid=True, squared_pred=True, reduction='mean')
+# Load the config file
+with open("./config.yaml", "r") as ymlfile:
+   config_file = yaml.load(ymlfile, Loader=yaml.Loader)
+rank_list = [2, 4, 6, 8, 16, 32, 64, 128, 256, 512]
+rank_loss = []
+# Load SAM model
+seg_loss = monai.losses.DiceCELoss(sigmoid=True, squared_pred=True, reduction='mean')
+with torch.no_grad():
+    for rank in rank_list:
+        sam = build_sam_vit_b(checkpoint=config_file["SAM"]["CHECKPOINT"])
+        #Create SAM LoRA
+        sam_lora = LoRA_sam(sam, rank)
+        sam_lora.load_lora_parameters(f"./lora_weights/lora_rank{rank}.safetensors")  
+        model = sam_lora.sam
+
+        # Process the dataset
+        processor = Samprocessor(model)
+        dataset = DatasetSegmentation(config_file, processor, mode="test")
+
+        # Create a dataloader
+        test_dataloader = DataLoader(dataset, batch_size=1, collate_fn=collate_fn)
+
+
+        # Set model to train and into the device
+        model.eval()
+        model.to(device)
+
+
+        total_score = []
+        for i, batch in enumerate(tqdm(test_dataloader)):
+
+            outputs = model(batched_input=batch,
+                multimask_output=False)
+
+            gt_mask_tensor = batch[0]["ground_truth_mask"].unsqueeze(0).unsqueeze(0) # We need to get the [B, C, H, W] starting from [H, W]
+            loss = seg_loss(outputs[0]["low_res_logits"], gt_mask_tensor.float().to(device))
+
+            total_score.append(loss.item())
+
+
+        print(f'Mean dice score: {mean(total_score)}')
+        rank_loss.append(mean(total_score))
+
+
+print("RANK LOSS :", rank_loss)
+
+width = 0.25  # the width of the bars
+multiplier = 0
+models_results= {"Rank 2": rank_loss[0], 
+                 "Rank 4": rank_loss[1], 
+                 "Rank 6": rank_loss[2],
+                 "Rank 8": rank_loss[3],
+                 "Rank 16": rank_loss[4],
+                 "Rank 32": rank_loss[5],
+                 "Rank 64": rank_loss[6],
+                 "Rank 128": rank_loss[7],
+                 "Rank 256": rank_loss[8],
+                 "Rank 512": rank_loss[9]
+                 }
+eval_scores_name = ["Rank"]
+x = np.arange(len(eval_scores_name))
+fig, ax = plt.subplots(layout='constrained')
+
+for model_name, score in models_results.items():
+    offset = width * multiplier
+    rects = ax.bar(x + offset, score, width, label=model_name)
+    ax.bar_label(rects, padding=3)
+    multiplier += 1
+
+# Add some text for labels, title and custom x-axis tick labels, etc.
+ax.set_ylabel('Dice Loss')
+ax.set_title('LoRA trained on 50 epochs - Rank comparison on test set')
+ax.set_xticks(x + width, eval_scores_name)
+ax.legend(loc=3, ncols=2)
+ax.set_ylim(0, 0.15)
+
+plt.savefig("./plots/rank_comparison.jpg")
diff --git a/inference.py → inference_plots.py b/inference.py → inference_plots.py
@@ -14,7 +14,7 @@
 sam_checkpoint = "sam_vit_b_01ec64.pth"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 sam = build_sam_vit_b(checkpoint=sam_checkpoint)
-rank = 64
+rank = 512
 sam_lora = LoRA_sam(sam, rank)
 sam_lora.load_lora_parameters(f"./lora_weights/lora_rank{rank}.safetensors")
 model = sam_lora.sam
@@ -25,7 +25,7 @@ def inference_model(sam_model, image_path, filename, mask_path=None, bbox=None,
         model = sam_model.sam
         rank = sam_model.rank
     else:
-        model = sam_model
+        model = build_sam_vit_b(checkpoint=sam_checkpoint)
 
     model.eval()
     model.to(device)
@@ -55,9 +55,10 @@ def inference_model(sam_model, image_path, filename, mask_path=None, bbox=None,
         ax2.imshow(masks[0])
         if is_baseline:
             ax2.set_title(f"Baseline SAM prediction: {filename}")
+            plt.savefig(f"./plots/{filename}_baseline.jpg")
         else:
             ax2.set_title(f"SAM LoRA rank {rank} prediction: {filename}")
-        plt.savefig("./plots/" + filename)
+            plt.savefig(f"./plots/{filename[:-4]}_rank{rank}.jpg")
 
     else:
         fig, (ax1, ax2, ax3) = plt.subplots(1, 3, sharex=True, sharey=True, figsize=(15, 15))
@@ -72,9 +73,10 @@ def inference_model(sam_model, image_path, filename, mask_path=None, bbox=None,
         ax3.imshow(masks[0])
         if is_baseline:
             ax3.set_title(f"Baseline SAM prediction: {filename}")
+            plt.savefig(f"./plots/{filename}_baseline.jpg")
         else:
             ax3.set_title(f"SAM LoRA rank {rank} prediction: {filename}")
-        plt.savefig("./plots/" + filename)
+            plt.savefig(f"./plots/{filename[:-4]}_rank{rank}.jpg")
 
 
 # Open configuration file
@@ -91,16 +93,16 @@ def inference_model(sam_model, image_path, filename, mask_path=None, bbox=None,
 inference_train = False
 
 if inference_train:
-    total_loss = []
+
     for image_name, dict_annot in train_set.items():
         image_path = f"./dataset/train/images/{image_name}"
-        inference_model(sam_lora, image_path, filename=image_name, mask_path=dict_annot["mask_path"], bbox=dict_annot["bbox"], is_baseline=False)
+        inference_model(sam_lora, image_path, filename=image_name, mask_path=dict_annot["mask_path"], bbox=dict_annot["bbox"], is_baseline=True)
 
 
 else:
-    total_loss = []
+
     for image_name, dict_annot in test_set.items():
         image_path = f"./dataset/test/images/{image_name}"
-        inference_model(sam_lora, image_path, filename=image_name, bbox=dict_annot["bbox"], is_baseline=False)
+        inference_model(sam_lora, image_path, filename=image_name, mask_path=dict_annot["mask_path"], bbox=dict_annot["bbox"], is_baseline=True)
 
 
diff --git a/plots/baseline/preds_test/ring_test_1.jpg b/plots/baseline/preds_test/ring_test_1.jpg
diff --git a/plots/baseline/preds_test/ring_test_1.jpg_baseline.jpg b/plots/baseline/preds_test/ring_test_1.jpg_baseline.jpg
diff --git a/plots/baseline/preds_test/ring_test_2.jpg b/plots/baseline/preds_test/ring_test_2.jpg
diff --git a/plots/baseline/preds_test/ring_test_2.jpg_baseline.jpg b/plots/baseline/preds_test/ring_test_2.jpg_baseline.jpg
diff --git a/plots/baseline/preds_test/ring_test_3.jpg b/plots/baseline/preds_test/ring_test_3.jpg
diff --git a/plots/baseline/preds_test/ring_test_4.jpg b/plots/baseline/preds_test/ring_test_4.jpg
diff --git a/plots/baseline/preds_test/ring_test_5.jpg b/plots/baseline/preds_test/ring_test_5.jpg
diff --git a/plots/baseline/preds_train/ring1.jpg b/plots/baseline/preds_train/ring1.jpg
diff --git a/plots/baseline/preds_train/ring1.jpg_baseline.jpg b/plots/baseline/preds_train/ring1.jpg_baseline.jpg
diff --git a/plots/baseline/preds_train/ring2.jpg → ...seline/preds_train/ring2.jpg_baseline.jpg b/plots/baseline/preds_train/ring2.jpg → ...seline/preds_train/ring2.jpg_baseline.jpg
diff --git a/plots/baseline/preds_train/ring3.jpg b/plots/baseline/preds_train/ring3.jpg
diff --git a/plots/baseline/preds_train/ring3.jpg_baseline.jpg b/plots/baseline/preds_train/ring3.jpg_baseline.jpg
diff --git a/plots/baseline/preds_train/ring4.jpg b/plots/baseline/preds_train/ring4.jpg
diff --git a/plots/baseline/preds_train/ring4.jpg_baseline.jpg b/plots/baseline/preds_train/ring4.jpg_baseline.jpg
diff --git a/plots/baseline/preds_train/ring5.jpg b/plots/baseline/preds_train/ring5.jpg
diff --git a/plots/baseline/preds_train/ring5.jpg_baseline.jpg b/plots/baseline/preds_train/ring5.jpg_baseline.jpg
diff --git a/plots/baseline/preds_train/ring6.jpg_baseline.jpg b/plots/baseline/preds_train/ring6.jpg_baseline.jpg
diff --git a/plots/baseline/preds_train/ring7.jpg_baseline.jpg b/plots/baseline/preds_train/ring7.jpg_baseline.jpg
diff --git a/plots/baseline/preds_train/ring8.jpg_baseline.jpg b/plots/baseline/preds_train/ring8.jpg_baseline.jpg
diff --git a/plots/best_model_rank_512/on_test/ring_test_1_rank512.jpg b/plots/best_model_rank_512/on_test/ring_test_1_rank512.jpg
diff --git a/plots/best_model_rank_512/on_test/ring_test_2_rank512.jpg b/plots/best_model_rank_512/on_test/ring_test_2_rank512.jpg
diff --git a/plots/best_model_rank_512/on_train/ring1_rank512.jpg b/plots/best_model_rank_512/on_train/ring1_rank512.jpg
diff --git a/plots/best_model_rank_512/on_train/ring2_rank512.jpg b/plots/best_model_rank_512/on_train/ring2_rank512.jpg
diff --git a/plots/best_model_rank_512/on_train/ring3_rank512.jpg b/plots/best_model_rank_512/on_train/ring3_rank512.jpg
diff --git a/plots/best_model_rank_512/on_train/ring4_rank512.jpg b/plots/best_model_rank_512/on_train/ring4_rank512.jpg
diff --git a/plots/best_model_rank_512/on_train/ring5_rank512.jpg b/plots/best_model_rank_512/on_train/ring5_rank512.jpg
diff --git a/plots/best_model_rank_512/on_train/ring6_rank512.jpg b/plots/best_model_rank_512/on_train/ring6_rank512.jpg
diff --git a/plots/best_model_rank_512/on_train/ring7_rank512.jpg b/plots/best_model_rank_512/on_train/ring7_rank512.jpg
diff --git a/plots/best_model_rank_512/on_train/ring8_rank512.jpg b/plots/best_model_rank_512/on_train/ring8_rank512.jpg
diff --git a/plots/rank_comparison.jpg b/plots/rank_comparison.jpg
diff --git a/plots/ring1.jpg b/plots/ring1.jpg
diff --git a/plots/ring2.jpg b/plots/ring2.jpg
diff --git a/plots/ring3.jpg b/plots/ring3.jpg
diff --git a/plots/ring4.jpg b/plots/ring4.jpg
diff --git a/plots/ring5.jpg b/plots/ring5.jpg
diff --git a/plots/ring_test_1.jpg b/plots/ring_test_1.jpg
diff --git a/plots/ring_test_2.jpg b/plots/ring_test_2.jpg
diff --git a/plots/ring_test_3.jpg b/plots/ring_test_3.jpg
diff --git a/src/dataloader.py b/src/dataloader.py
@@ -40,13 +40,6 @@ def __init__(self, config_file: dict, processor: Samprocessor, mode: str):
             for img_path in self.img_files:
                 self.mask_files.append(os.path.join(config_file["DATASET"]["TRAIN_PATH"],'masks', os.path.basename(img_path)[:-4] + ".jpg")) 
 
-        elif mode == "valid":
-
-            self.img_files = glob.glob(os.path.join(config_file["DATASET"]["VALID_PATH"],'images','*.jpg'))
-            self.mask_files = []
-            for img_path in self.img_files:
-                self.mask_files.append(os.path.join(config_file["DATASET"]["VALID_PATH"],'masks', os.path.basename(img_path)[:-4] + ".jpg"))
-
         else:
             self.img_files = glob.glob(os.path.join(config_file["DATASET"]["TEST_PATH"],'images','*.jpg'))
             self.mask_files = []