fix: bug when saving grey level images introduced in 12ebb72

LaboratoireMecaniqueLille · Oct 9, 2024 · 118df0b · 118df0b
1 parent 12ebb72
commit 118df0b
Showing 1 changed file with 13 additions and 5 deletions.
diff --git a/src/crappy/blocks/camera_processes/record.py b/src/crappy/blocks/camera_processes/record.py
@@ -226,13 +226,21 @@ def loop(self) -> None:
     # Saving the image at the destination path using the chosen backend
     self.log(logging.DEBUG, "Saving image")
     if self._save_backend == 'sitk':
-      Sitk.WriteImage(Sitk.GetImageFromArray(self.img[:, :, ::-1],
-                                             isVector=True), path)
+      if len(self.img.shape) == 3:
+        Sitk.WriteImage(Sitk.GetImageFromArray(self.img[:, :, ::-1],
+                                               isVector=True), path)
+      else:
+        Sitk.WriteImage(Sitk.GetImageFromArray(self.img), path)
 
     elif self._save_backend == 'pil':
-      PIL.Image.fromarray(self.img[:, :, ::-1]).save(
-        path, exif={TAGS_INV[key]: val for key, val in self.metadata.items()
-                    if key in TAGS_INV})
+      if len(self.img.shape) == 3:
+        PIL.Image.fromarray(self.img[:, :, ::-1]).save(
+          path, exif={TAGS_INV[key]: val for key, val in self.metadata.items()
+                      if key in TAGS_INV})
+      else:
+        PIL.Image.fromarray(self.img).save(
+          path, exif={TAGS_INV[key]: val for key, val in self.metadata.items()
+                      if key in TAGS_INV})
 
     elif self._save_backend == 'cv2':
       cv2.imwrite(path, self.img)