mu map dataset now return torch tensors with channel dimension

mu_map/data/datasets.py

@@ -5,6 +5,7 @@ import cv2 as cv
 import pandas as pd
 import pydicom
 import numpy as np
+import torch
 from torch.utils.data import Dataset
 
 from mu_map.data.prepare import headers
@@ -86,12 +87,20 @@ class MuMapDataset(Dataset):
                 bed_contour = self.bed_contours[row["id"]]
                 for i in range(mu_map.shape[0]):
                     mu_map[i] = cv.drawContours(mu_map[i], [bed_contour], -1, 0.0, -1)
-            self.mu_maps[_id] = mu_map
 
             recon_file = os.path.join(self.dir_images, row[headers.file_recon_nac_nsc])
             recon = pydicom.dcmread(recon_file).pixel_array
             if self.align:
                 recon = align_images(recon, mu_map)
+
+            mu_map = mu_map.astype(np.float32)
+            mu_map = torch.from_numpy(mu_map)
+            mu_map = mu_map.unsqueeze(dim=0)
+            self.mu_maps[_id] = mu_map
+
+            recon = recon.astype(np.float32)
+            recon = torch.from_numpy(recon)
+            recon = recon.unsqueeze(dim=0)
             self.reconstructions[_id] = recon
         print("Pre-loading images done!")
 
@@ -194,6 +203,8 @@ if __name__ == "__main__":
         im = 0
 
         recon, mu_map = dataset[i]
+        recon = recon.squeeze().numpy()
+        mu_map = mu_map.squeeze().numpy()
         print(f"{(i+1):>{len(str(len(dataset)))}}/{len(dataset)}", end="\r")
 
         cv.imshow(wname, combine_images((recon, mu_map), (ir, im)))
@@ -206,8 +217,6 @@ if __name__ == "__main__":
 
             to_show = combine_images((recon, mu_map), (ir, im))
             cv.imshow(wname, to_show)
-
-
             key = cv.waitKey(timeout)
 
             if key == ord("n"):