shinestacker 1.7.0__py3-none-any.whl → 1.8.1__py3-none-any.whl

shinestacker/algorithms/stack_framework.py

@@ -46,7 +46,8 @@ class StackJob(Job):
 class ImageSequenceManager:
     def __init__(self, name, input_path='', output_path='', working_path='',
                  plot_path=constants.DEFAULT_PLOTS_PATH,
-                 scratch_output_dir=True, resample=1,
+                 scratch_output_dir=True, delete_output_at_end=False,
+                 resample=1,
                  reverse_order=constants.DEFAULT_FILE_REVERSE_ORDER, **_kwargs):
         self.name = name
         self.working_path = working_path
@@ -56,6 +57,7 @@ class ImageSequenceManager:
         self._resample = resample
         self.reverse_order = reverse_order
         self.scratch_output_dir = scratch_output_dir
+        self.delete_output_at_end = delete_output_at_end
         self.enabled = None
         self.base_message = ''
         self._input_full_path = None
@@ -122,6 +124,24 @@ class ImageSequenceManager:
                                      constants.LOG_COLOR_LEVEL_2))
         self.base_message = color_str(self.name, constants.LOG_COLOR_LEVEL_1, "bold")
 
+    def scratch_outout_folder(self):
+        if self.enabled:
+            output_dir = self.output_full_path()
+            list_dir = os.listdir(output_dir)
+            n_files = len(list_dir)
+            if n_files > 0:
+                for filename in list_dir:
+                    file_path = os.path.join(output_dir, filename)
+                    if os.path.isfile(file_path):
+                        os.remove(file_path)
+            self.print_message(
+                color_str(f"output directory {self.output_path} content erased",
+                          'yellow'))
+        else:
+            self.print_message(
+                color_str(f"module disabled, output directory {self.output_path}"
+                          " not scratched", 'yellow'))
+
     def init(self, job):
         if self.working_path == '':
             self.working_path = job.working_path
@@ -130,22 +150,10 @@ class ImageSequenceManager:
         if not os.path.exists(output_dir):
             os.makedirs(output_dir)
         else:
-            list_dir = os.listdir(output_dir)
-            if len(list_dir) > 0:
+            if len(os.listdir(output_dir)):
                 if self.scratch_output_dir:
-                    if self.enabled:
-                        for filename in list_dir:
-                            file_path = os.path.join(output_dir, filename)
-                            if os.path.isfile(file_path):
-                                os.remove(file_path)
-                        self.print_message(
-                            color_str(f": output directory {self.output_path} content erased",
-                                      'yellow'))
-                    else:
-                        self.print_message(
-                            color_str(f": module disabled, output directory {self.output_path}"
-                                      " not scratched", 'yellow'))
-                else:
+                    self.scratch_outout_folder()
+                elif self.enabled:
                     self.print_message(
                         color_str(
                             f": output directory {self.output_path} not empty, "
@@ -168,6 +176,11 @@ class ImageSequenceManager:
                     self._input_filepaths.append(filepath)
         job.add_action_path(self.output_path)
 
+    def end_job(self):
+        if self.delete_output_at_end:
+            self.scratch_outout_folder()
+            os.rmdir(self.output_full_path())
+
     def folder_list_str(self):
         if isinstance(self.input_full_path(), list):
             file_list = ", ".join(
@@ -178,7 +191,8 @@ class ImageSequenceManager:
 
 
 class ReferenceFrameTask(SequentialTask, ImageSequenceManager):
-    def __init__(self, name, enabled=True, reference_index=0, step_process=False, **kwargs):
+    def __init__(self, name, enabled=True, reference_index=0,
+                 step_process=constants.DEFAULT_COMBINED_ACTIONS_STEP_PROCESS, **kwargs):
         ImageSequenceManager.__init__(self, name, **kwargs)
         SequentialTask.__init__(self, name, enabled)
         self.ref_idx = reference_index
@@ -206,6 +220,9 @@ class ReferenceFrameTask(SequentialTask, ImageSequenceManager):
     def end(self):
         SequentialTask.end(self)
 
+    def end_job(self):
+        ImageSequenceManager.end_job(self)
+
     def run_frame(self, _idx, _ref_idx):
         return None
 
@@ -260,7 +277,8 @@ class SubAction:
 
 class CombinedActions(ReferenceFrameTask):
     def __init__(self, name, actions=[], enabled=True, **kwargs):
-        ReferenceFrameTask.__init__(self, name, enabled, **kwargs)
+        step_process = kwargs.pop('step_process', constants.DEFAULT_COMBINED_ACTIONS_STEP_PROCESS)
+        ReferenceFrameTask.__init__(self, name, enabled, step_process=step_process, **kwargs)
         self._actions = actions
         self._metadata = (None, None)
 
@@ -278,11 +296,15 @@ class CombinedActions(ReferenceFrameTask):
         self._metadata = get_img_metadata(img)
         return img
 
+    def frame_str(self, idx=-1):
+        if self.run_sequential():
+            idx = self.current_action_count
+        return f"frame {idx + 1}/{self.total_action_counts}"
+
     def run_frame(self, idx, ref_idx):
         input_path = self.input_filepath(idx)
         self.print_message(
-            color_str(f'read input image '
-                      f'{idx + 1}/{self.total_action_counts}, '
+            color_str(f'read input {self.frame_str(idx)}, '
                       f'{os.path.basename(input_path)}', constants.LOG_COLOR_LEVEL_3))
         img = read_img(input_path)
         validate_image(img, *(self._metadata))
@@ -301,20 +323,19 @@ class CombinedActions(ReferenceFrameTask):
                 if img is not None:
                     img = a.run_frame(idx, ref_idx, img)
                 else:
-                    self.sub_message(
-                        color_str(": null input received, action skipped",
+                    self.print_message(
+                        color_str("null input received, action skipped",
                                   constants.LOG_COLOR_ALERT),
                         level=logging.WARNING)
         if img is not None:
             output_path = os.path.join(self.output_full_path(), os.path.basename(input_path))
             self.print_message(
-                color_str(f'write output image '
-                          f'{idx + 1}/{self.total_action_counts}, '
+                color_str(f'write output {self.frame_str(idx)}, '
                           f'{os.path.basename(output_path)}', constants.LOG_COLOR_LEVEL_3))
             write_img(output_path, img)
             return img
         self.print_message(color_str(
-            f"no output file resulted from processing input file: {os.path.basename(input_path)}",
+            f"no output resulted from processing input file: {os.path.basename(input_path)}",
             constants.LOG_COLOR_ALERT), level=logging.WARNING)
         return None
 
@@ -323,6 +344,9 @@ class CombinedActions(ReferenceFrameTask):
             if a.enabled:
                 a.end()
 
+    def end_job(self):
+        ReferenceFrameTask.end_job(self)
+
     def sequential_processing(self):
         for a in self._actions:
             if a.sequential_processing():

shinestacker/algorithms/utils.py

@@ -140,8 +140,12 @@ def save_plot(filename, fig=None):
         if not os.path.isdir(dir_path):
             os.makedirs(dir_path)
         if fig is None:
+            logging_level = logging.getLogger().level
+            logger = logging.getLogger()
+            logger.setLevel(logging.WARNING)
             fig = plt.gcf()
-        fig.savefig(filename, dpi=150)
+            fig.savefig(filename, dpi=150)
+            logger.setLevel(logging_level)
         if config.JUPYTER_NOTEBOOK:
             plt.show()
         plt.close(fig)

shinestacker/algorithms/vignetting.py

@@ -4,7 +4,7 @@ import traceback
 import logging
 import numpy as np
 import matplotlib.pyplot as plt
-from scipy.optimize import curve_fit, fsolve
+from scipy.optimize import curve_fit, bisect
 import cv2
 from .. core.colors import color_str
 from .. core.core_utils import setup_matplotlib_mode
@@ -181,9 +181,22 @@ class Vignetting(SubAction):
             self.process.callback(
                 constants.CALLBACK_SAVE_PLOT, self.process.id,
                 f"{self.process.name}: intensity\nframe {idx_str}", plot_path)
+
         for i, p in enumerate(self.percentiles):
-            self.corrections[i][idx] = fsolve(lambda x: sigmoid_model(x, *params) /
-                                              self.v0 - p, r0_fit)[0]
+            s1 = sigmoid_model(0, *params) / self.v0
+            s2 = sigmoid_model(self.r_max, *params) / self.v0
+            if s1 > p and s2 < p:
+                try:
+                    c = bisect(lambda x: sigmoid_model(x, *params) / self.v0 - p, 0, self.r_max)
+                except Exception as e:
+                    traceback.print_tb(e.__traceback__)
+                    self.process.sub_message(color_str(f": {str(e).lower()}", "yellow"),
+                                             level=logging.WARNING)
+            elif s1 <= p:
+                c = 0
+            else:
+                c = self.r_max
+            self.corrections[i][idx] = c
         self.process.print_message(
             color_str(f"{self.process.idx_tot_str(idx)}: correct vignetting", "cyan"))
         return correct_vignetting(