Fixed output filename formatting which assumed the extension was alwa…

…ys 3 characters long (#43)
teamtomo · Dec 14, 2023 · 1f9747a · 1f9747a
1 parent d76e97e
commit 1f9747a
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/src/membrain_seg/segmentation/dataloading/data_utils.py b/src/membrain_seg/segmentation/dataloading/data_utils.py
@@ -175,7 +175,8 @@ def store_segmented_tomograms(
     out_folder = out_folder
     if store_probabilities:
         out_file = os.path.join(
-            out_folder, os.path.basename(orig_data_path)[:-4] + "_scores.mrc"
+            out_folder,
+            os.path.splitext(os.path.basename(orig_data_path))[0] + "_scores.mrc",
         )
         out_tomo = Tomogram(
             data=predictions_np, header=mrc_header, voxel_size=voxel_size
@@ -186,7 +187,10 @@ def store_segmented_tomograms(
     )
     out_file_thres = os.path.join(
         out_folder,
-        os.path.basename(orig_data_path)[:-4] + "_" + ckpt_token + "_segmented.mrc",
+        os.path.splitext(os.path.basename(orig_data_path))[0]
+        + "_"
+        + ckpt_token
+        + "_segmented.mrc",
     )
     if store_connected_components:
         predictions_np_thres = connected_components(