From 4d245893d5a54e912a2ff305dcf74766a0926cf4 Mon Sep 17 00:00:00 2001
From: Tamino Huxohl <thuxohl@techfak.uni-bielefeld.de>
Date: Wed, 21 Dec 2022 09:46:05 +0100
Subject: [PATCH] slight changes to test script for visulization
---
mu_map/test.py | 22 ++++++++++++++--------
1 file changed, 14 insertions(+), 8 deletions(-)
diff --git a/mu_map/test.py b/mu_map/test.py
index a11ec3c..f761307 100644
--- a/mu_map/test.py
+++ b/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:
--
GitLab