Implement generalization to allow for custom data

Author: Jewelry-x

cmp_test_label.py

@@ -3,14 +3,11 @@
 from model import UNet
 import torch
 import os
-
-P_label = "result\PTest_test_label.npy"
-T_label = "result\TTest_test_label.npy"
-model_path = "result\saved_model_pt"
+from config import *
 
 # Check if the model file exists
-if os.path.exists(model_path):
-    whole_model = torch.load(model_path)
+if os.path.exists(os.path.join(RESULT_PATH, "saved_model_pt")):
+    whole_model = torch.load(os.path.join(RESULT_PATH, "saved_model_pt"))
     # Construct the model details string
     model_details = ""
     if "data" in whole_model:
@@ -19,20 +16,17 @@
         model_details += (
             "Learning rate: " + str(whole_model.get("learning_rate")) + "\n"
         )
-    if "phantom_IOU" in whole_model:
-        model_details += "Phantom IoU: " + str(whole_model.get("phantom_IOU")) + "\n"
-    if "phantom_DC" in whole_model:
-        model_details += "Phantom DC: " + str(whole_model.get("phantom_DC")) + "\n"
-    if "T1T6_IOU" in whole_model:
-        model_details += "T1T6 IoU: " + str(whole_model.get("T1T6_IOU")) + "\n"
-    if "T1T6_DC" in whole_model:
-        model_details += "T1T6 DC: " + str(whole_model.get("T1T6_DC")) + "\n"
     if "pools" in whole_model:
         model_details += "Pooling layers used: " + str(whole_model.get("pools")) + "\n"
+    if "reverse_pools" in whole_model:
+        model_details += "Reverse pooling used: " + str(whole_model.get("reverse_pools")) + "\n"
     if "data_augmentations" in whole_model:
         model_details += (
             "Data Augmentations: " + str(whole_model.get("data_augmentations")) + "\n"
         )
+    if "tested_on" in whole_model:
+        for i, data in enumerate(whole_model["tested_on"]):
+            model_details += f"{data} IOU: {whole_model['IOU'][i]}\n{data} DC: {whole_model['DC'][i]}\n"
     if "image_size" in whole_model:
         model_details += "Image Size: " + str(whole_model.get("image_size")) + "\n"
     if "optimizer" in whole_model:
@@ -52,49 +46,38 @@
 else:
     model_details = "Model not found or loaded"
 
-# Create subplots for images and model details
-fig, axs = plt.subplots(2, 4, gridspec_kw={"width_ratios": [1, 1, 1, 0.5]})
-
-# Check if P_label exists
-if os.path.exists(P_label):
-    Phantom_created = np.load(P_label)
-    Pimage = np.load("data\\PTest\\frame_0000.npy")
-    Pmask = np.load("data\\PTest_label\\frame_0000.npy")
 
-    # Plot Phantom images
-    axs[0, 0].imshow(np.squeeze(Pimage), cmap="gray")
-    axs[0, 0].set_title("Phantom Original Image")
+# Get the list of files in the directory
+all_files = os.listdir(RESULT_PATH)
 
-    axs[0, 1].imshow(np.squeeze(Pmask), cmap="gray")
-    axs[0, 1].set_title("Phantom Mask")
+# Filter the files based on the extension
+filtered_files = [file for file in all_files if file.endswith(".npy")]
 
-    axs[0, 2].imshow(np.squeeze(Phantom_created), cmap="gray")
-    axs[0, 2].set_title("Phantom Model Created Mask")
+# Create subplots for images and model details
+fig, axs = plt.subplots(len(filtered_files), 4, gridspec_kw={"width_ratios": [1, 1, 1, 0.5]})
 
-# Check if T_label exists
-if os.path.exists(T_label):
-    T1T6_created = np.load(T_label)
-    Timage = np.load("data\\TTest\\frame_0000.npy")
-    Tmask = np.load("data\\TTest_label\\frame_0000.npy")
+idx = 0
+for file in filtered_files:
+    data = file[:-15]
+    created_mask = np.load(os.path.join(RESULT_PATH, file))
+    original_mask = np.load(os.path.join(DATA_PATH, TESTING_DATA_MASK_LOCATION[TESTING_DATA.index(data)], MASK_DEFINITION % 0000))
+    original_image = np.load(os.path.join(DATA_PATH, TESTING_DATA_LOCATION[TESTING_DATA.index(data)], IMAGE_DEFINITION % 0000) )
 
-    # Plot T1-T6 images
-    axs[1, 0].imshow(np.squeeze(Timage), cmap="gray")
-    axs[1, 0].set_title("T1-T6 Original Image")
+    # Plot Phantom images
+    axs[idx, 0].imshow(np.squeeze(original_image), cmap="gray")
+    axs[idx, 0].set_title(data + " Original Image")
 
-    axs[1, 1].imshow(np.squeeze(Tmask), cmap="gray")
-    axs[1, 1].set_title("T1-T6 Mask")
+    axs[idx, 1].imshow(np.squeeze(original_mask), cmap="gray")
+    axs[idx, 1].set_title(data + " Mask")
 
-    axs[1, 2].imshow(np.squeeze(T1T6_created), cmap="gray")
-    axs[1, 2].set_title("T1-T6 Model Created Mask")
+    axs[idx, 2].imshow(np.squeeze(created_mask), cmap="gray")
+    axs[idx, 2].set_title(data + " Model Created Mask")
 
-# Hide ticks and labels for the empty subplot
-axs[0, 3].axis("off")
+    axs[idx, 3].axis("off")
 
-# Adjust layout to add space for the text
-plt.subplots_adjust(bottom=0.1, top=0.9)
+    idx += 1
 
-# Add model details below the images
+# Hide ticks and labels for the empty subplot
 axs[0, 3].text(0, 0.5, model_details, fontsize=14, ha="left", va="center")
-axs[1, 3].axis("off")
 
-plt.show()
+plt.show()

config.py

@@ -1,6 +1,6 @@
-DATA = "P"
+DATA = "Phantom"
 TRAIN = True
-TEST_ON_BOTH_DATA = False
+TEST_ON = [0, 1]
 POOL = 4
 REVERSE_POOL = True
 LEARNING_RATE = 1e-2
@@ -17,12 +17,29 @@
 CREATE_TEST_MASK = True
 
 RESULT_PATH = "./result"
-DATA_PATH = "data"
 
 CREATE_FOLDER = True
 
 NOTIFY = True
 
+# DATA LOADING
+DATA_PATH = "data"
+
+TRAINING_DATA = ["Phantom", "T1T6"]
+TESTING_DATA = ["Phantom", "T1T6"]
+
+TRAINING_DATA_LOCATION = ["PTrain", "TTrain"]
+TESTING_DATA_LOCATION = ["PTest", "TTest"]
+
+TRAINING_DATA_MASK_LOCATION = ["PTrain_label", "TTrain_label"]
+TESTING_DATA_MASK_LOCATION = ["PTest_label", "TTest_label"]
+
+IMAGE_DEFINITION = "frame_%04d.npy"
+MASK_DEFINITION = "frame_%04d.npy"
+
+TRAINING_DATA_COUNT = [1400, 845]
+TESTING_DATA_COUNT = [600, 362]
+
 # Transformations
 TRANSFORM = False
 

graph_progress.py

@@ -1,9 +1,11 @@
 import openpyxl
 import matplotlib.pyplot as plt
 import mplcursors
+import os
+from config import RESULT_PATH
 
 # Load the workbook
-workbook = openpyxl.load_workbook("./result/progress.xlsx")
+workbook = openpyxl.load_workbook(os.path.join(RESULT_PATH, "progress.xlsx"))
 worksheet = workbook.active
 
 # Initialize lists to store data