Fixed error when No Stretch was selected

Steffenhir · Apr 2, 2024 · 0f70695 · 0f70695
1 parent 5b84ff3
commit 0f70695
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 19 deletions.
diff --git a/graxpert/AstroImageRepository.py b/graxpert/AstroImageRepository.py
@@ -16,30 +16,39 @@ def stretch_all(self, stretch_params:StretchParameters, saturation:float):
         if self.get("Original") is None:
             return
 
-        all_image_arrays = []
-        all_mtf_stretch_params = []
+        stretches = []
 
-        all_image_arrays.append(self.get("Original").img_array)
-        all_mtf_stretch_params.append(calculate_mtf_stretch_parameters_for_image(stretch_params, self.get("Original").img_array))
-
-        if self.get("Gradient-Corrected") is not None and self.get("Background") is not None:
-            all_image_arrays.append(self.get("Gradient-Corrected").img_array)
-            all_mtf_stretch_params.append(calculate_mtf_stretch_parameters_for_image(stretch_params, self.get("Gradient-Corrected").img_array))
+        if not stretch_params.do_stretch:
+            for key, image in self.images.items():
+                if image is not None:
+                    stretches.append(image.img_array)
+
+        else:
+
+            all_image_arrays = []
+            all_mtf_stretch_params = []
 
-            all_image_arrays.append(self.get("Background").img_array)
-            all_mtf_stretch_params.append(all_mtf_stretch_params[0])
+            all_image_arrays.append(self.get("Original").img_array)
+            all_mtf_stretch_params.append(calculate_mtf_stretch_parameters_for_image(stretch_params, self.get("Original").img_array))
 
-
-        if self.get("Denoised") is not None and self.get("Gradient-Corrected") is None:
-            all_image_arrays.append(self.get("Denoised").img_array)
-            all_mtf_stretch_params.append(all_mtf_stretch_params[0])
+            if self.get("Gradient-Corrected") is not None and self.get("Background") is not None:
+                all_image_arrays.append(self.get("Gradient-Corrected").img_array)
+                all_mtf_stretch_params.append(calculate_mtf_stretch_parameters_for_image(stretch_params, self.get("Gradient-Corrected").img_array))
+
+                all_image_arrays.append(self.get("Background").img_array)
+                all_mtf_stretch_params.append(all_mtf_stretch_params[0])
+
 
-        elif self.get("Denoised") is not None and self.get("Gradient-Corrected") is not None:
-            all_image_arrays.append(self.get("Denoised").img_array)
-            all_mtf_stretch_params.append(all_mtf_stretch_params[1])
+            if self.get("Denoised") is not None and self.get("Gradient-Corrected") is None:
+                all_image_arrays.append(self.get("Denoised").img_array)
+                all_mtf_stretch_params.append(all_mtf_stretch_params[0])
+
+            elif self.get("Denoised") is not None and self.get("Gradient-Corrected") is not None:
+                all_image_arrays.append(self.get("Denoised").img_array)
+                all_mtf_stretch_params.append(all_mtf_stretch_params[1])
+
 
-
-        stretches = stretch_all(all_image_arrays, all_mtf_stretch_params)
+            stretches = stretch_all(all_image_arrays, all_mtf_stretch_params)
 
 
         i = 0

diff --git a/graxpert/stretch.py b/graxpert/stretch.py
@@ -53,6 +53,9 @@ def __init__(self, stretch_option: str, channels_linked: bool = False, images_li
 
 
 def stretch(data, stretch_params: StretchParameters):
+    if not stretch_params.do_stretch:
+        return data
+
     mtf_stretch_param = calculate_mtf_stretch_parameters_for_image(stretch_params, data)
     return stretch_all([data], [mtf_stretch_param])[0]