shinestacker 1.0.4.post2__py3-none-any.whl → 1.2.0__py3-none-any.whl

shinestacker/algorithms/base_stack_algo.py

@@ -1,9 +1,10 @@
 # pylint: disable=C0114, C0115, C0116, E0602, R0903
+import os
 import numpy as np
-from .. core.exceptions import InvalidOptionError, ImageLoadError
+from .. core.exceptions import InvalidOptionError, ImageLoadError, RunStopException
 from .. config.constants import constants
 from .. core.colors import color_str
-from .utils import read_img, get_img_metadata, validate_image
+from .utils import read_img, get_img_metadata, validate_image, get_img_file_shape, extension_tif_jpg
 
 
 class BaseStackAlgo:
@@ -11,6 +12,9 @@ class BaseStackAlgo:
         self._name = name
         self._steps_per_frame = steps_per_frame
         self.process = None
+        self.filenames = None
+        self.shape = None
+        self.do_step_callback = False
         if float_type == constants.FLOAT_32:
             self.float_type = np.float32
         elif float_type == constants.FLOAT_64:
@@ -24,8 +28,23 @@ class BaseStackAlgo:
     def name(self):
         return self._name
 
-    def steps_per_frame(self):
-        return self._steps_per_frame
+    def set_process(self, process):
+        self.process = process
+
+    def set_do_step_callback(self, enable):
+        self.do_step_callback = enable
+
+    def init(self, filenames):
+        self.filenames = filenames
+        first_img_file = ''
+        for filename in filenames:
+            if os.path.isfile(filename) and extension_tif_jpg(filename):
+                first_img_file = filename
+                break
+        self.shape = get_img_file_shape(first_img_file)
+
+    def total_steps(self, n_frames):
+        return self._steps_per_frame * n_frames
 
     def print_message(self, msg):
         self.process.sub_message_r(color_str(msg, constants.LOG_COLOR_LEVEL_3))
@@ -40,3 +59,13 @@ class BaseStackAlgo:
         else:
             validate_image(img, *metadata)
         return img, metadata, updated
+
+    def check_running(self, cleanup_callback=None):
+        if self.process.callback('check_running', self.process.id, self.process.name) is False:
+            if cleanup_callback is not None:
+                cleanup_callback()
+            raise RunStopException(self.name)
+
+    def after_step(self, step):
+        if self.do_step_callback:
+            self.process.callback('after_step', self.process.id, self.process.name, step)

shinestacker/algorithms/depth_map.py

@@ -2,7 +2,7 @@
 import numpy as np
 import cv2
 from .. config.constants import constants
-from .. core.exceptions import InvalidOptionError, RunStopException
+from .. core.exceptions import InvalidOptionError
 from .utils import read_img, img_bw
 from .base_stack_algo import BaseStackAlgo
 
@@ -61,19 +61,18 @@ class DepthMapStack(BaseStackAlgo):
         raise InvalidOptionError("map_type", self.map_type, details=f" valid values are "
                                  f"{constants.DM_MAP_AVERAGE} and {constants.DM_MAP_MAX}.")
 
-    def focus_stack(self, filenames):
+    def focus_stack(self):
         gray_images = []
         metadata = None
-        for i, img_path in enumerate(filenames):
+        for i, img_path in enumerate(self.filenames):
             self.print_message(f": reading file (1/2) {img_path.split('/')[-1]}")
 
             img, metadata, _updated = self.read_image_and_update_metadata(img_path, metadata)
 
             gray = img_bw(img)
             gray_images.append(gray)
-            self.process.callback('after_step', self.process.id, self.process.name, i)
-            if self.process.callback('check_running', self.process.id, self.process.name) is False:
-                raise RunStopException(self.name)
+            self.after_step(i)
+            self.check_running()
         dtype = metadata[1]
         gray_images = np.array(gray_images, dtype=self.float_type)
         if self.energy == constants.DM_ENERGY_SOBEL:
@@ -92,7 +91,7 @@ class DepthMapStack(BaseStackAlgo):
             energies = self.smooth_energy(energies)
         weights = self.get_focus_map(energies)
         blended_pyramid = None
-        for i, img_path in enumerate(filenames):
+        for i, img_path in enumerate(self.filenames):
             self.print_message(f": reading file (2/2) {img_path.split('/')[-1]}")
             img = read_img(img_path).astype(self.float_type)
             weight = weights[i]
@@ -110,14 +109,12 @@ class DepthMapStack(BaseStackAlgo):
                              for j in range(self.levels)]
             blended_pyramid = current_blend if blended_pyramid is None \
                 else [np.add(bp, cb) for bp, cb in zip(blended_pyramid, current_blend)]
-            self.process.callback('after_step', self.process.id,
-                                  self.process.name, i + len(filenames))
-            if self.process.callback('check_running', self.process.id, self.process.name) is False:
-                raise RunStopException(self.name)
+            self.after_step(i + len(self.filenames))
+            self.check_running()
         result = blended_pyramid[0]
         self.print_message(': blend levels')
         for j in range(1, self.levels):
             size = (blended_pyramid[j].shape[1], blended_pyramid[j].shape[0])
             result = cv2.pyrUp(result, dstsize=size) + blended_pyramid[j]
-        n_values = 255 if dtype == np.uint8 else 65535
+        n_values = constants.MAX_UINT8 if dtype == np.uint8 else constants.MAX_UINT16
         return np.clip(np.absolute(result), 0, n_values).astype(dtype)

shinestacker/algorithms/multilayer.py

@@ -1,4 +1,4 @@
-# pylint: disable=C0114, C0115, C0116, E1101, R0914, E0606
+# pylint: disable=C0114, C0115, C0116, E1101, R0914, E0606, R0912
 import os
 import logging
 import cv2
@@ -61,6 +61,13 @@ def write_multilayer_tiff_from_images(image_dict, output_file, exif_path='', cal
     if len(dtypes) > 1:
         raise RuntimeError("All input files must all have 8 bit or 16 bit depth.")
     dtype = dtypes[0]
+    bytes_per_pixel = 3 * np.dtype(dtype).itemsize
+    est_memory = shape[0] * shape[1] * bytes_per_pixel * len(image_dict)
+    if est_memory > constants.MULTILAYER_WARNING_MEM_GB * constants.ONE_GIGA:
+        if callbacks:
+            callback = callbacks.get('memory_warning', None)
+            if callback:
+                callback(float(est_memory) / constants.ONE_GIGA)
     max_pixel_value = constants.MAX_UINT16 if dtype == np.uint16 else constants.MAX_UINT8
     transp = np.full_like(list(image_dict.values())[0][..., 0], max_pixel_value)
     compression_type = PsdCompressionType.ZIP_PREDICTED
@@ -203,7 +210,10 @@ class MultiLayer(JobBase, FrameMultiDirectory):
             'exif_msg': lambda path: self.print_message(
                 color_str(f"copying exif data from path: {path}", constants.LOG_COLOR_LEVEL_2)),
             'write_msg': lambda path: self.print_message(
-                color_str(f"writing multilayer tiff file: {path}", constants.LOG_COLOR_LEVEL_2))
+                color_str(f"writing multilayer tiff file: {path}", constants.LOG_COLOR_LEVEL_2)),
+            'memory_warning': lambda mem: self.print_message(
+                color_str(f"warning: estimated file size: {mem:.2f} GBytes",
+                          constants.LOG_COLOR_WARNING))
         }
         write_multilayer_tiff(input_files, output_file, labels=None, exif_path=self.exif_path,
                               callbacks=callbacks)

shinestacker/algorithms/noise_detection.py

@@ -1,4 +1,4 @@
-# pylint: disable=C0114, C0115, C0116, E1101, W0718, R0914, R0915
+# pylint: disable=C0114, C0115, C0116, E1101, W0718, R0914, R0915, R0902
 import os
 import errno
 import logging
@@ -11,7 +11,7 @@ from .. core.colors import color_str
 from .. core.exceptions import ImageLoadError
 from .. core.framework import JobBase
 from .. core.core_utils import make_tqdm_bar
-from .. core.exceptions import RunStopException
+from .. core.exceptions import RunStopException, ShapeError
 from .stack_framework import FrameMultiDirectory, SubAction
 from .utils import read_img, save_plot, get_img_metadata, validate_image
 
@@ -49,7 +49,7 @@ class NoiseDetection(JobBase, FrameMultiDirectory):
     def __init__(self, name="noise-map", enabled=True, **kwargs):
         FrameMultiDirectory.__init__(self, name, **kwargs)
         JobBase.__init__(self, name, enabled)
-        self.max_frames = kwargs.get('max_frames', -1)
+        self.max_frames = kwargs.get('max_frames', constants.DEFAULT_NOISE_MAX_FRAMES)
         self.blur_size = kwargs.get('blur_size', constants.DEFAULT_BLUR_SIZE)
         self.file_name = kwargs.get('file_name', constants.DEFAULT_NOISE_MAP_FILENAME)
         if self.file_name == '':
@@ -154,6 +154,7 @@ class MaskNoise(SubAction):
         self.method = method
         self.process = None
         self.noise_mask_img = None
+        self.expected_shape = None
 
     def begin(self, process):
         self.process = process
@@ -163,6 +164,7 @@ class MaskNoise(SubAction):
                 f': reading noisy pixel mask file: {self.noise_mask}',
                 constants.LOG_COLOR_LEVEL_3))
             self.noise_mask_img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
+            self.expected_shape = self.noise_mask_img.shape[:2]
             if self.noise_mask_img is None:
                 raise ImageLoadError(path, f"failed to load image file {self.noise_mask}.")
         else:
@@ -170,6 +172,9 @@ class MaskNoise(SubAction):
 
     def run_frame(self, _idx, _ref_idx, image):
         self.process.sub_message_r(color_str(': mask noisy pixels', constants.LOG_COLOR_LEVEL_3))
+        shape = image.shape[:2]
+        if shape != self.expected_shape:
+            raise ShapeError(self.expected_shape, shape)
         if len(image.shape) == 3:
             corrected = image.copy()
             for c in range(3):

shinestacker/algorithms/pyramid.py

@@ -1,25 +1,36 @@
-# pylint: disable=C0114, C0115, C0116, E1101
+# pylint: disable=C0114, C0115, C0116, E1101, R0913, R0917, R0902
 import numpy as np
 import cv2
 from .. config.constants import constants
-from .. core.exceptions import RunStopException
 from .utils import read_img
 from .base_stack_algo import BaseStackAlgo
 
 
 class PyramidBase(BaseStackAlgo):
-    def __init__(self, min_size=constants.DEFAULT_PY_MIN_SIZE,
+    def __init__(self, name, min_size=constants.DEFAULT_PY_MIN_SIZE,
                  kernel_size=constants.DEFAULT_PY_KERNEL_SIZE,
                  gen_kernel=constants.DEFAULT_PY_GEN_KERNEL,
                  float_type=constants.DEFAULT_PY_FLOAT):
-        super().__init__("pyramid", 2, float_type)
+        super().__init__(name, 2, float_type)
         self.min_size = min_size
         self.kernel_size = kernel_size
         self.pad_amount = (kernel_size - 1) // 2
-        self.do_step_callback = False
         kernel = np.array([0.25 - gen_kernel / 2.0, 0.25,
                            gen_kernel, 0.25, 0.25 - gen_kernel / 2.0])
         self.gen_kernel = np.outer(kernel, kernel)
+        self.dtype = None
+        self.num_pixel_values = None
+        self.max_pixel_value = None
+        self.n_levels = 0
+        self.n_frames = 0
+
+    def init(self, filenames):
+        super().init(filenames)
+        self.n_levels = int(np.log2(min(self.shape) / self.min_size))
+
+    def total_steps(self, n_frames):
+        self.n_frames = n_frames
+        return self._steps_per_frame * n_frames + self.n_levels
 
     def convolve(self, image):
         return cv2.filter2D(image, -1, self.gen_kernel, borderType=cv2.BORDER_REFLECT101)
@@ -55,6 +66,7 @@ class PyramidBase(BaseStackAlgo):
         return fused
 
     def collapse(self, pyramid):
+        self.print_message(': collapsing pyramid')
         img = pyramid[-1]
         for layer in pyramid[-2::-1]:
             expanded = self.expand_layer(img)
@@ -110,19 +122,23 @@ class PyramidBase(BaseStackAlgo):
             fused += np.where(best_d[:, :, np.newaxis] == layer, img, 0)
         return (fused / 2).astype(images.dtype)
 
+    def focus_stack_validate(self, cleanup_callback=None):
+        metadata = None
+        n = len(self.filenames)
+        for i, img_path in enumerate(self.filenames):
+            self.print_message(f": validating file {img_path.split('/')[-1]}, {i + 1}/{n}")
 
-class PyramidStack(PyramidBase):
-    def __init__(self, min_size=constants.DEFAULT_PY_MIN_SIZE,
-                 kernel_size=constants.DEFAULT_PY_KERNEL_SIZE,
-                 gen_kernel=constants.DEFAULT_PY_GEN_KERNEL,
-                 float_type=constants.DEFAULT_PY_FLOAT):
-        super().__init__(min_size, kernel_size, gen_kernel, float_type)
-        self.offset = np.arange(-self.pad_amount, self.pad_amount + 1)
-        self.dtype = None
-        self.num_pixel_values = None
-        self.max_pixel_value = None
+            _img, metadata, updated = self.read_image_and_update_metadata(img_path, metadata)
+            if updated:
+                self.dtype = metadata[1]
+                self.num_pixel_values = constants.NUM_UINT8 \
+                    if self.dtype == np.uint8 else constants.NUM_UINT16
+                self.max_pixel_value = constants.MAX_UINT8 \
+                    if self.dtype == np.uint8 else constants.MAX_UINT16
+            self.after_step(i + 1)
+            self.check_running(cleanup_callback)
 
-    def process_single_image(self, img, levels):
+    def single_image_laplacian(self, img, levels):
         pyramid = [img.astype(self.float_type)]
         for _ in range(levels):
             next_layer = self.reduce_layer(pyramid[-1])
@@ -136,44 +152,43 @@ class PyramidStack(PyramidBase):
             h, w = pyr.shape[:2]
             expanded = expanded[:h, :w]
             laplacian.append(pyr - expanded)
+        return laplacian
+
+
+class PyramidStack(PyramidBase):
+    def __init__(self, min_size=constants.DEFAULT_PY_MIN_SIZE,
+                 kernel_size=constants.DEFAULT_PY_KERNEL_SIZE,
+                 gen_kernel=constants.DEFAULT_PY_GEN_KERNEL,
+                 float_type=constants.DEFAULT_PY_FLOAT):
+        super().__init__("pyramid", min_size, kernel_size, gen_kernel, float_type)
+        self.offset = np.arange(-self.pad_amount, self.pad_amount + 1)
+
+    def process_single_image(self, img, levels):
+        laplacian = self.single_image_laplacian(img, levels)
         return laplacian[::-1]
 
     def fuse_pyramids(self, all_laplacians):
         fused = [self.get_fused_base(np.stack([p[-1] for p in all_laplacians], axis=0))]
+        count = 0
         for layer in range(len(all_laplacians[0]) - 2, -1, -1):
             self.print_message(f': fusing pyramids, layer: {layer + 1}')
             laplacians = np.stack([p[layer] for p in all_laplacians], axis=0)
             fused.append(self.fuse_laplacian(laplacians))
+            count += 1
+            self.after_step(self._steps_per_frame * self.n_frames + count)
+            self.check_running()
         self.print_message(': pyramids fusion completed')
         return fused[::-1]
 
-    def focus_stack(self, filenames):
-        metadata = None
+    def focus_stack(self):
+        n = len(self.filenames)
+        self.focus_stack_validate()
         all_laplacians = []
-        levels = None
-        n = len(filenames)
-        for i, img_path in enumerate(filenames):
-            self.print_message(f": validating file {img_path.split('/')[-1]}")
-
-            img, metadata, updated = self.read_image_and_update_metadata(img_path, metadata)
-            if updated:
-                self.dtype = metadata[1]
-                self.num_pixel_values = constants.NUM_UINT8 \
-                    if self.dtype == np.uint8 else constants.NUM_UINT16
-                self.max_pixel_value = constants.MAX_UINT8 \
-                    if self.dtype == np.uint8 else constants.MAX_UINT16
-                levels = int(np.log2(min(img.shape[:2]) / self.min_size))
-            if self.do_step_callback:
-                self.process.callback('after_step', self.process.id, self.process.name, i)
-            if self.process.callback('check_running', self.process.id, self.process.name) is False:
-                raise RunStopException(self.name)
-        for i, img_path in enumerate(filenames):
-            self.print_message(f": processing file {img_path.split('/')[-1]}")
+        for i, img_path in enumerate(self.filenames):
+            self.print_message(f": processing file {img_path.split('/')[-1]} ({i + 1}/{n})")
             img = read_img(img_path)
-            all_laplacians.append(self.process_single_image(img, levels))
-            if self.do_step_callback:
-                self.process.callback('after_step', self.process.id, self.process.name, i + n)
-            if self.process.callback('check_running', self.process.id, self.process.name) is False:
-                raise RunStopException(self.name)
+            all_laplacians.append(self.process_single_image(img, self.n_levels))
+            self.after_step(i + n + 1)
+            self.check_running()
         stacked_image = self.collapse(self.fuse_pyramids(all_laplacians))
         return stacked_image.astype(self.dtype)

shinestacker/algorithms/pyramid_auto.py

@@ -0,0 +1,141 @@
+# pylint: disable=C0114, C0115, C0116, E1101, R0913, R0902, R0914, R0917
+import os
+import numpy as np
+from .. config.constants import constants
+from .utils import extension_tif_jpg
+from .base_stack_algo import BaseStackAlgo
+from .pyramid import PyramidStack
+from .pyramid_tiles import PyramidTilesStack
+
+
+class PyramidAutoStack(BaseStackAlgo):
+    def __init__(self, min_size=constants.DEFAULT_PY_MIN_SIZE,
+                 kernel_size=constants.DEFAULT_PY_KERNEL_SIZE,
+                 gen_kernel=constants.DEFAULT_PY_GEN_KERNEL,
+                 float_type=constants.DEFAULT_PY_FLOAT,
+                 tile_size=constants.DEFAULT_PY_TILE_SIZE,
+                 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,
+                 mode='auto'):
+        super().__init__("auto_pyramid", 2, float_type)
+        self.min_size = min_size
+        self.kernel_size = kernel_size
+        self.gen_kernel = gen_kernel
+        self.float_type = float_type
+        self.tile_size = tile_size
+        self.n_tiled_layers = n_tiled_layers
+        self.memory_limit = memory_limit * constants.ONE_GIGA
+        self.max_threads = max_threads
+        available_cores = os.cpu_count() or 1
+        self.num_threads = min(max_threads, available_cores)
+        self.max_tile_size = max_tile_size
+        self.min_n_tiled_layers = min_n_tiled_layers
+        self.mode = mode
+        self._implementation = None
+        self.dtype = None
+        self.shape = None
+        self.n_levels = None
+        self.n_frames = 0
+        self.channels = 3
+        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
+
+    def init(self, filenames):
+        first_img_file = None
+        for filename in filenames:
+            if os.path.isfile(filename) and extension_tif_jpg(filename):
+                first_img_file = filename
+                break
+        if first_img_file is None:
+            raise ValueError("No valid image files found")
+        _img, metadata, _ = self.read_image_and_update_metadata(first_img_file, None)
+        self.shape, self.dtype = metadata
+        self.n_levels = int(np.log2(min(self.shape) / self.min_size))
+        self.n_frames = len(filenames)
+        memory_required_memory = self._estimate_memory_memory()
+        if self.mode == 'memory' or (self.mode == 'auto' and
+                                     memory_required_memory <= self.memory_limit):
+            self._implementation = PyramidStack(
+                min_size=self.min_size,
+                kernel_size=self.kernel_size,
+                gen_kernel=self.gen_kernel,
+                float_type=self.float_type
+            )
+            self.print_message(": using memory-based pyramid stacking")
+        else:
+            optimal_params = self._find_optimal_tile_params()
+            self._implementation = PyramidTilesStack(
+                min_size=self.min_size,
+                kernel_size=self.kernel_size,
+                gen_kernel=self.gen_kernel,
+                float_type=self.float_type,
+                tile_size=optimal_params['tile_size'],
+                n_tiled_layers=optimal_params['n_tiled_layers'],
+                max_threads=self.num_threads
+            )
+            self.print_message(f": using tile-based pyramid stacking "
+                               f"(tile_size: {optimal_params['tile_size']}, "
+                               f"n_tiled_layers: {optimal_params['n_tiled_layers']}), "
+                               f"{self.num_threads} cores.")
+        self._implementation.init(filenames)
+        self._implementation.set_do_step_callback(self.do_step_callback)
+        if self.process is not None:
+            self._implementation.set_process(self.process)
+        else:
+            raise RuntimeError("self.process must be initialized.")
+
+    def _estimate_memory_memory(self):
+        h, w = self.shape[:2]
+        total_memory = 0
+        for _ in range(self.n_levels):
+            total_memory += h * w * self.bytes_per_pixel
+            h, w = max(1, h // 2), max(1, w // 2)
+        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)))
+        tile_size = min(self.max_tile_size, tile_size_max, self.shape[0], self.shape[1])
+        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:
+                n_tiled_layers = layer + 1
+            else:
+                break
+        n_tiled_layers = max(n_tiled_layers, self.min_n_tiled_layers)
+        n_tiled_layers = min(n_tiled_layers, self.n_levels)
+        return {'tile_size': tile_size, 'n_tiled_layers': n_tiled_layers}
+
+    def set_process(self, process):
+        super().set_process(process)
+        if self._implementation is not None:
+            self._implementation.set_process(process)
+
+    def total_steps(self, n_frames):
+        if self._implementation is None:
+            return super().total_steps(n_frames)
+        return self._implementation.total_steps(n_frames)
+
+    def focus_stack(self):
+        if self._implementation is None:
+            raise RuntimeError("PyramidAutoStack not initialized")
+        return self._implementation.focus_stack()
+
+    def after_step(self, step):
+        if self._implementation is not None:
+            self._implementation.after_step(step)
+        else:
+            super().after_step(step)
+
+    def check_running(self, cleanup_callback=None):
+        if self._implementation is not None:
+            self._implementation.check_running(cleanup_callback)
+        else:
+            super().check_running(cleanup_callback)