Add fix for uncertainty training

radionets-project · Apr 19, 2024 · 608cc86 · 608cc86
1 parent 33460e0
commit 608cc86
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/radionets/dl_framework/callbacks.py b/radionets/dl_framework/callbacks.py
@@ -66,6 +66,7 @@ def plot_test_pred(self):
         pred = rescale_normalization(pred, norm_dict)
         if pred.shape[1] == 4:
             self.uncertainty = True
+            pred = torch.stack((pred[:, 0, :], pred[:, 2, :]), dim=1)
         images = {"pred": pred, "truth": img_true}
         images = apply_symmetry(images)
         pred = images["pred"]
@@ -74,10 +75,7 @@ def plot_test_pred(self):
         fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(16, 10))
         lim_phase = check_vmin_vmax(img_true[0, 1])
         im1 = ax1.imshow(pred[0, 0], cmap="inferno")
-        if self.uncertainty:
-            im2 = ax2.imshow(pred[0, 2], cmap=OrBu, vmin=-lim_phase, vmax=lim_phase)
-        else:
-            im2 = ax2.imshow(pred[0, 1], cmap=OrBu, vmin=-lim_phase, vmax=lim_phase)
+        im2 = ax2.imshow(pred[0, 1], cmap=OrBu, vmin=-lim_phase, vmax=lim_phase)
         im3 = ax3.imshow(img_true[0, 0], cmap="inferno")
         im4 = ax4.imshow(img_true[0, 1], cmap=OrBu, vmin=-lim_phase, vmax=lim_phase)
         make_axes_nice(fig, ax1, im1, "Real")
@@ -103,6 +101,8 @@ def plot_test_fft(self):
             with torch.no_grad():
                 pred = eval_model(img_test, model)
         pred = rescale_normalization(pred, norm_dict)
+        if self.uncertainty:
+            pred = torch.stack((pred[:, 0, :], pred[:, 2, :]), dim=1)
         images = {"pred": pred, "truth": img_true}
         images = apply_symmetry(images)
         pred = images["pred"]