slight changes to test script for visulization

mu_map/test.py

@@ -5,29 +5,31 @@ import torch
 
 from mu_map.dataset.default import MuMapDataset
 from mu_map.dataset.mock import MuMapMockDataset
-from mu_map.dataset.normalization import norm_max
+from mu_map.dataset.normalization import norm_max, norm_gaussian
 from mu_map.models.unet import UNet
 from mu_map.util import to_grayscale, COLOR_WHITE
 
 torch.set_grad_enabled(False)
 
-dataset = MuMapMockDataset("data/initial/")
+dataset = MuMapMockDataset("data/second/")
 
-model = UNet(in_channels=1, features=[8, 16])
+# model = UNet(in_channels=1, features=[8, 16])
+model = UNet(in_channels=1)
 device = torch.device("cpu")
-weights = torch.load("train_data/snapshots/10.pth", map_location=device)
+weights = torch.load("train_data/snapshots/val_min_Model.pth", map_location=device)
 model.load_state_dict(weights)
 model = model.eval()
 
 recon, mu_map = dataset[0]
 recon = recon.unsqueeze(dim=0)
-recon = norm_max(recon)
+# recon = norm_max(recon)
+recon = norm_gaussian(recon)
 
 output = model(recon)
 # output = output * 40206.0
 
 diff = ((mu_map - output) ** 2).mean()
-print(f"Diff: {diff:.3f}")
+print(f"Diff: {diff:.5f}")
 
 output = output.squeeze().numpy()
 mu_map = mu_map.squeeze().numpy()
@@ -39,8 +41,11 @@ cv.namedWindow(wname, cv.WINDOW_NORMAL)
 cv.resizeWindow(wname, 1600, 900)
 space = np.full((1024, 10), 239, np.uint8)
 
-def to_display_image(image, _slice):
-    _image = to_grayscale(image[_slice], min_val=image.min(), max_val=image.max())
+def to_display_image(image, _slice, _min=None, _max=None):
+    _max = _max if _max is not None else image.max()
+    _min = _min if _min is not None else image.min()
+
+    _image = to_grayscale(image[_slice], min_val=_min, max_val=_max)
     _image = cv.resize(_image, (1024, 1024), cv.INTER_AREA)
     _text = f"{str(_slice):>{len(str(image.shape[0]))}}/{str(image.shape[0])}"
     _image = cv.putText(
@@ -54,6 +59,7 @@ def com(image1, image2, _slice):
     space = np.full((image1.shape[0], 10), 239, np.uint8)
     return np.hstack((image1, space, image2))
 
+output = np.clip(output, 0, mu_map.max())
 
 i = 0
 while True: