shinestacker 1.2.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

shinestacker/algorithms/pyramid_auto.py

@@ -17,8 +17,9 @@ class PyramidAutoStack(BaseStackAlgo):
                  n_tiled_layers=constants.DEFAULT_PY_N_TILED_LAYERS,
                  memory_limit=constants.DEFAULT_PY_MEMORY_LIMIT_GB,
                  max_threads=constants.DEFAULT_PY_MAX_THREADS,
-                 max_tile_size=2048,
-                 min_n_tiled_layers=1,
+                 max_tile_size=constants.DEFAULT_PY_MAX_TILE_SIZE,
+                 min_tile_size=constants.DEFAULT_PY_MIN_TILE_SIZE,
+                 min_n_tiled_layers=constants.DEFAULT_PY_MIN_N_TILED_LAYERS,
                  mode='auto'):
         super().__init__("auto_pyramid", 2, float_type)
         self.min_size = min_size
@@ -32,6 +33,7 @@ class PyramidAutoStack(BaseStackAlgo):
         available_cores = os.cpu_count() or 1
         self.num_threads = min(max_threads, available_cores)
         self.max_tile_size = max_tile_size
+        self.min_tile_size = min_tile_size
         self.min_n_tiled_layers = min_n_tiled_layers
         self.mode = mode
         self._implementation = None
@@ -39,10 +41,10 @@ class PyramidAutoStack(BaseStackAlgo):
         self.shape = None
         self.n_levels = None
         self.n_frames = 0
-        self.channels = 3
+        self.channels = 3  # r, g, b
         dtype = np.float32 if self.float_type == constants.FLOAT_32 else np.float64
         self.bytes_per_pixel = self.channels * np.dtype(dtype).itemsize
-        self.overhead = 1.5
+        self.overhead = constants.PY_MEMORY_OVERHEAD
 
     def init(self, filenames):
         first_img_file = None
@@ -97,15 +99,19 @@ class PyramidAutoStack(BaseStackAlgo):
         return self.overhead * total_memory * self.n_frames
 
     def _find_optimal_tile_params(self):
-        tile_size_max = int(np.sqrt(self.memory_limit /
-                            (self.num_threads * self.n_frames *
-                             self.bytes_per_pixel * self.overhead)))
+        h, w = self.shape[:2]
+        base_level_memory = h * w * self.bytes_per_pixel
+        available_memory = self.memory_limit - base_level_memory
+        available_memory /= self.overhead
+        tile_size_max = int(np.sqrt(available_memory /
+                            (self.num_threads * self.n_frames * self.bytes_per_pixel)))
         tile_size = min(self.max_tile_size, tile_size_max, self.shape[0], self.shape[1])
+        tile_size = max(self.min_tile_size, tile_size)
         n_tiled_layers = 0
         for layer in range(self.n_levels):
-            h = max(1, self.shape[0] // (2 ** layer))
-            w = max(1, self.shape[1] // (2 ** layer))
-            if h > tile_size or w > tile_size:
+            h_layer = max(1, self.shape[0] // (2 ** layer))
+            w_layer = max(1, self.shape[1] // (2 ** layer))
+            if h_layer > tile_size or w_layer > tile_size:
                 n_tiled_layers = layer + 1
             else:
                 break

shinestacker/algorithms/pyramid_tiles.py

@@ -2,10 +2,11 @@
 # pylint: disable=C0114, C0115, C0116, E1101, R0914, R1702, R1732, R0913
 # pylint: disable=R0917, R0912, R0915, R0902, W0718
 import os
+import gc
 import time
 import shutil
 import tempfile
-import concurrent.futures
+from concurrent.futures import ThreadPoolExecutor, as_completed
 import numpy as np
 from .. config.constants import constants
 from .. core.exceptions import RunStopException
@@ -121,13 +122,13 @@ class PyramidTilesStack(PyramidBase):
             for x in range(0, w, self.tile_size):
                 tiles.append((y, x))
         self.print_message(f': starting parallel propcessging on {self.num_threads} cores')
-        with concurrent.futures.ThreadPoolExecutor(max_workers=self.num_threads) as executor:
+        with ThreadPoolExecutor(max_workers=self.num_threads) as executor:
             future_to_tile = {
                 executor.submit(
                     self._process_tile, level, num_images, all_level_counts, y, x, h, w): (y, x)
                 for y, x in tiles
             }
-            for future in concurrent.futures.as_completed(future_to_tile):
+            for future in as_completed(future_to_tile):
                 y, x = future_to_tile[future]
                 try:
                     fused_tile = future.result()
@@ -154,7 +155,9 @@ class PyramidTilesStack(PyramidBase):
         if laplacians:
             stacked = np.stack(laplacians, axis=0)
             return self.fuse_laplacian(stacked)
-        y_end, x_end = min(y + self.tile_size, h), min(x + self.tile_size, w)
+        y_end = min(y + self.tile_size, h)
+        x_end = min(x + self.tile_size, w)
+        gc.collect()
         return np.zeros((y_end - y, x_end - x, 3), dtype=self.float_type)
 
     def fuse_pyramids(self, all_level_counts, num_images):
@@ -202,26 +205,29 @@ class PyramidTilesStack(PyramidBase):
         self.focus_stack_validate(self.cleanup_temp_files)
         all_level_counts = [0] * n
         if self.num_threads > 1:
-            self.print_message(f': starting parallel image processing on {self.num_threads} cores')
+            self.print_message(f': starting parallel processing on {self.num_threads} cores')
             args_list = [(file_path, i, n) for i, file_path in enumerate(self.filenames)]
             executor = None
             try:
-                executor = concurrent.futures.ThreadPoolExecutor(max_workers=self.num_threads)
+                executor = ThreadPoolExecutor(max_workers=self.num_threads)
                 future_to_index = {
                     executor.submit(self._process_single_image_wrapper, args): i
                     for i, args in enumerate(args_list)
                 }
                 completed_count = 0
-                for future in concurrent.futures.as_completed(future_to_index):
+                for future in as_completed(future_to_index):
                     i = future_to_index[future]
                     try:
                         img_index, level_count = future.result()
                         all_level_counts[img_index] = level_count
                         completed_count += 1
-                        self.print_message(f': completed processing image {completed_count}/{n}')
+                        self.print_message(
+                            ": processing completed, image  "
+                            f"{self.idx_tot_str(completed_count - 1)}")
                     except Exception as e:
-                        self.print_message(f"Error processing image {i + 1}: {str(e)}")
-                    self.after_step(i + n + 1)
+                        self.print_message(
+                            f"Error processing image  {self.idx_tot_str(i)}: {str(e)}")
+                    self.after_step(completed_count + n + 1)
                     self.check_running(lambda: None)
             except RunStopException:
                 self.print_message(": stopping image processing...")
@@ -235,7 +241,9 @@ class PyramidTilesStack(PyramidBase):
                     executor.shutdown(wait=True)
         else:
             for i, file_path in enumerate(self.filenames):
-                self.print_message(f": processing file {file_path.split('/')[-1]}, {i + 1}/{n}")
+                self.print_message(
+                    f": processing file {os.path.basename(file_path)}, "
+                    f"{self.idx_tot_str(i)}")
                 img = read_img(file_path)
                 level_count = self.process_single_image(img, self.n_levels, i)
                 all_level_counts[i] = level_count

shinestacker/algorithms/stack.py

@@ -2,19 +2,19 @@
 import os
 import numpy as np
 from .. config.constants import constants
-from .. core.framework import JobBase
+from .. core.framework import TaskBase
 from .. core.colors import color_str
 from .. core.exceptions import InvalidOptionError
 from .utils import write_img, extension_tif_jpg
-from .stack_framework import FrameDirectory, ActionList
+from .stack_framework import ImageSequenceManager, SequentialTask
 from .exif import copy_exif_from_file_to_file
 from .denoise import denoise
 
 
-class FocusStackBase(JobBase, FrameDirectory):
+class FocusStackBase(TaskBase, ImageSequenceManager):
     def __init__(self, name, stack_algo, enabled=True, **kwargs):
-        FrameDirectory.__init__(self, name, **kwargs)
-        JobBase.__init__(self, name, enabled)
+        ImageSequenceManager.__init__(self, name, **kwargs)
+        TaskBase.__init__(self, name, enabled)
         self.stack_algo = stack_algo
         self.exif_path = kwargs.pop('exif_path', '')
         self.prefix = kwargs.pop('prefix', constants.DEFAULT_STACK_PREFIX)
@@ -26,9 +26,10 @@ class FocusStackBase(JobBase, FrameDirectory):
     def focus_stack(self, filenames):
         self.sub_message_r(color_str(': reading input files', constants.LOG_COLOR_LEVEL_3))
         stacked_img = self.stack_algo.focus_stack()
-        in_filename = filenames[0].split(".")
-        out_filename = f"{self.output_dir}/{self.prefix}{in_filename[0]}." + \
-            '.'.join(in_filename[1:])
+        in_filename = os.path.basename(filenames[0]).split(".")
+        out_filename = os.path.join(
+            self.output_full_path(),
+            f"{self.prefix}{in_filename[0]}." + '.'.join(in_filename[1:]))
         if self.denoise_amount > 0:
             self.sub_message_r(': denoise image')
             stacked_img = denoise(stacked_img, self.denoise_amount, self.denoise_amount)
@@ -45,14 +46,14 @@ class FocusStackBase(JobBase, FrameDirectory):
             name = f"{self.name}: {self.stack_algo.name()}"
             if idx_str != '':
                 name += f"\nbunch: {idx_str}"
-            self.callback('save_plot', self.id, name, out_filename)
+            self.callback(constants.CALLBACK_SAVE_PLOT, self.id, name, out_filename)
         if self.frame_count >= 0:
             self.frame_count += 1
 
     def init(self, job, working_path=''):
-        FrameDirectory.init(self, job)
+        ImageSequenceManager.init(self, job)
         if self.exif_path is None:
-            self.exif_path = job.paths[0]
+            self.exif_path = job.action_path(0)
         if self.exif_path != '':
             self.exif_path = working_path + "/" + self.exif_path
 
@@ -63,9 +64,9 @@ def get_bunches(collection, n_frames, n_overlap):
     return bunches
 
 
-class FocusStackBunch(ActionList, FocusStackBase):
+class FocusStackBunch(SequentialTask, FocusStackBase):
     def __init__(self, name, stack_algo, enabled=True, **kwargs):
-        ActionList.__init__(self, name, enabled)
+        SequentialTask.__init__(self, name, enabled)
         FocusStackBase.__init__(self, name, stack_algo, enabled, **kwargs)
         self._chunks = None
         self.frame_count = 0
@@ -77,25 +78,28 @@ class FocusStackBunch(ActionList, FocusStackBase):
             raise InvalidOptionError("overlap", self.overlap,
                                      "overlap must be smaller than batch size")
 
+    def sequential_processing(self):
+        return True
+
     def init(self, job, _working_path=''):
         FocusStackBase.init(self, job, self.working_path)
 
     def begin(self):
-        ActionList.begin(self)
-        fnames = self.folder_filelist()
-        self._chunks = get_bunches(fnames, self.frames, self.overlap)
+        SequentialTask.begin(self)
+        self._chunks = get_bunches(self.input_filepaths(), self.frames, self.overlap)
         self.set_counts(len(self._chunks))
 
     def end(self):
-        ActionList.end(self)
+        SequentialTask.end(self)
 
-    def run_step(self):
-        self.print_message_r(color_str(f"fusing bunch: {self.count + 1}/{self.counts}",
-                                       constants.LOG_COLOR_LEVEL_2))
-        img_files = [os.path.join(self.input_full_path, name)
-                     for name in self._chunks[self.count - 1]]
+    def run_step(self, action_count=-1):
+        self.print_message(
+            color_str(f"fusing bunch: {action_count + 1}/{self.total_action_counts}",
+                      constants.LOG_COLOR_LEVEL_2))
+        img_files = self._chunks[action_count - 1]
         self.stack_algo.init(img_files)
-        self.focus_stack(self._chunks[self.count - 1])
+        self.focus_stack(self._chunks[action_count - 1])
+        return True
 
 
 class FocusStack(FocusStackBase):
@@ -106,11 +110,11 @@ class FocusStack(FocusStackBase):
 
     def run_core(self):
         self.set_filelist()
-        img_files = sorted([os.path.join(self.input_full_path, name) for name in self.filenames])
+        img_files = sorted(self.input_filepaths())
         self.stack_algo.init(img_files)
         self.callback('step_counts', self.id, self.name,
-                      self.stack_algo.total_steps(len(self.filenames)))
-        self.focus_stack(self.filenames)
+                      self.stack_algo.total_steps(self.num_input_filepaths()))
+        self.focus_stack(img_files)
 
     def init(self, job, _working_path=''):
         FocusStackBase.init(self, job, self.working_path)