default dataset can now deal with missing contours

mu_map/dataset/default.py

@@ -89,12 +89,16 @@ class MuMapDataset(Dataset):
         self.table = pd.read_csv(self.csv_file)
         if split_name:
             self.table = split_csv(self.table, self.split_file)[split_name]
-        self.table["id"] = self.table["id"].apply(int)
+        self.table[headers.id] = self.table[headers.id].apply(int)
 
         self.discard_mu_map_slices = discard_mu_map_slices
         self.align = align
         self.scatter_correction = scatter_correction
-        self.header_recon = headers.file_recon_nac_sc if self.scatter_correction else headers.file_recon_nac_nsc
+        self.header_recon = (
+            headers.file_recon_nac_sc
+            if self.scatter_correction
+            else headers.file_recon_nac_nsc
+        )
 
         self.reconstructions = {}
         self.mu_maps = {}
@@ -104,7 +108,7 @@ class MuMapDataset(Dataset):
         self.logger.debug("Pre-loading images ...")
         for i in range(len(self.table)):
             row = self.table.iloc[i]
-            _id = row["id"]
+            _id = row[headers.id]
 
             mu_map_file = os.path.join(self.dir_images, row[headers.file_mu_map])
             mu_map = pydicom.dcmread(mu_map_file)
@@ -112,9 +116,14 @@ class MuMapDataset(Dataset):
             if self.discard_mu_map_slices:
                 mu_map = discard_slices(row, mu_map)
             if self.bed_contours:
-                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)
+                if _id in self.bed_contours:
+                    bed_contour = self.bed_contours[_id]
+                    for i in range(mu_map.shape[0]):
+                        mu_map[i] = cv.drawContours(
+                            mu_map[i], [bed_contour], -1, 0.0, -1
+                        )
+                else:
+                    logger.warning(f"Could not find bed contour for id {_id}")
 
             recon_file = os.path.join(self.dir_images, row[self.header_recon])
             recon = pydicom.dcmread(recon_file)
@@ -138,8 +147,10 @@ class MuMapDataset(Dataset):
 
     def __getitem__(self, index: int):
         row = self.table.iloc[index]
-        _id = row["id"]
+        _id = row[headers.id]
+        return self.getitem_by_id(_id)
 
+    def getitem_by_id(self, _id: int):
         recon = self.reconstructions[_id]
         mu_map = self.mu_maps[_id]
 
@@ -153,6 +164,7 @@ class MuMapDataset(Dataset):
 
 __all__ = [MuMapDataset.__name__]
 
+
 def main(dataset):
     from mu_map.util import to_grayscale, COLOR_WHITE
 
@@ -182,10 +194,13 @@ def main(dataset):
         ir = 0
         im = 0
 
+        row = dataset.table.iloc[i]
+        _id = row[headers.id]
+
         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")
+        print(f"{(i+1):>{len(str(len(dataset)))}}/{len(dataset)} - ID: {_id}", end="\r")
 
         cv.imshow(wname, combine_images((recon, mu_map), (ir, im)))
         key = cv.waitKey(timeout)
@@ -265,4 +280,3 @@ if __name__ == "__main__":
         logger=logger,
     )
     main(dataset)
-