diff --git a/docs/api/core.md b/docs/api/core.md index 40373c4..6d97efb 100644 --- a/docs/api/core.md +++ b/docs/api/core.md @@ -178,11 +178,10 @@ at registration, all executors are assumed to be supported. {func}`~pyflowreg.core.backend_registry.register_backend`, {func}`~pyflowreg.core.backend_registry.get_backend`, -{func}`~pyflowreg.core.backend_registry.list_backends`, and -{func}`~pyflowreg.core.backend_registry.is_backend_available` are re-exported -from `pyflowreg.core`; -{func}`~pyflowreg.core.backend_registry.get_backend_executors` is available -from `pyflowreg.core.backend_registry`. +{func}`~pyflowreg.core.backend_registry.list_backends`, +{func}`~pyflowreg.core.backend_registry.is_backend_available`, and +{func}`~pyflowreg.core.backend_registry.get_backend_executors` are +re-exported from `pyflowreg.core`. ```{eval-rst} .. automodule:: pyflowreg.core.backend_registry @@ -216,8 +215,10 @@ flow (`cv2.DISOpticalFlow`) behind the same callable interface as {func}`~pyflowreg.core.optical_flow.get_displacement`. It is registered only when OpenCV (`cv2`) is importable. {class}`~pyflowreg.core.diso_optical_flow.DisoOF` reduces multi-channel input to grayscale using the channel weights, accepts -an optional initial flow field for warm starts, and initializes the OpenCV -DIS object lazily so instances remain picklable. +an optional initial flow field for warm starts (both as its native `w` +keyword and as the `uv` keyword the batch pipelines and `FlowRegLive` pass, +with `w` taking precedence), and initializes the OpenCV DIS object lazily so +instances remain picklable. Compared to the variational `flowreg` backend, `diso` has the following restrictions (enforced when the backend is resolved through `OFOptions`): diff --git a/docs/snippets/user_guide/online_processing/streaming_loop.py b/docs/snippets/user_guide/online_processing/streaming_loop.py index 60dd3bd..3806ea2 100644 --- a/docs/snippets/user_guide/online_processing/streaming_loop.py +++ b/docs/snippets/user_guide/online_processing/streaming_loop.py @@ -13,10 +13,10 @@ video = reader[:] reader.close() -# Scale to [0, 1]: incoming frames are normalized against the stored -# reference range, so they should arrive in a comparable intensity range +# Frames are normalized against the raw reference's intensity range +# internally, so they can stay at the recording's native scale (uint16 +# here); cast to float32 only to keep the demo arithmetic simple video = video.astype(np.float32) -video = (video - video.min()) / (video.max() - video.min()) # Configure optical flow; quality_setting is forced to "fast" internally options = OFOptions( diff --git a/docs/user_guide/backends.md b/docs/user_guide/backends.md index d548bdd..3916d63 100644 --- a/docs/user_guide/backends.md +++ b/docs/user_guide/backends.md @@ -28,6 +28,10 @@ Notes: resolved from `OFOptions`. - `diso` excludes multiprocessing because the multiprocessing workers import the variational solver directly and do not reconstruct registry backends. +- `diso` honors the flow initialization (`uv`) the pipelines pass between + batches and frames, forwarding it to OpenCV DIS as the initial flow; the + other variational solver keywords (`alpha`, `iterations`, ...) are accepted + but ignored. You can inspect what is registered in the current environment: diff --git a/docs/user_guide/configuration.md b/docs/user_guide/configuration.md index 116111b..8f1614a 100644 --- a/docs/user_guide/configuration.md +++ b/docs/user_guide/configuration.md @@ -194,7 +194,7 @@ The `flowreg_torch` backend additionally accepts `dtype` (`"float32"` or `"float :end-before: "[docs:end]" ``` -When a list of indices is given, the corresponding frames are read from the input, motion-compensated against their mean using increased regularization, and averaged to form the reference. +When a list of indices is given, the corresponding frames are read from the input, motion-compensated against their mean using increased regularization, and averaged to form the reference. Indices beyond the recording length are clipped to the last frame and the resulting duplicate indices are removed with a printed warning, so a short recording preregisters each available frame exactly once. When [cross-correlation pre-alignment](prealignment.md) is enabled, it also applies during this preregistration. ### Updating the Reference Frame diff --git a/docs/user_guide/file_formats.md b/docs/user_guide/file_formats.md index c9d5ea0..248c200 100644 --- a/docs/user_guide/file_formats.md +++ b/docs/user_guide/file_formats.md @@ -28,7 +28,7 @@ Readers are selected by file extension: `.tif`/`.tiff` (TIFF), `.h5`/`.hdf5`/`.h - Dataset names: `ch1`, `ch2`, ..., `chN` (pattern configurable via `dataset_names`) - File-level attributes: `frame_count`, `height`, `width`, `n_channels`, `dimension_ordering`, `format`, `dataset_names` -The reader interprets 3D datasets as `(T, H, W)` with one dataset per channel, and a 4D dataset as `(T, H, W, C)`. +The reader interprets 3D datasets as `(T, H, W)` with one dataset per channel by default, and a 4D dataset as `(T, H, W, C)`. For files stored in a different axis order, pass the same `dimension_ordering` option the writer uses (axis positions of height, width, time — e.g. `dimension_ordering=(0, 1, 2)` for `(H, W, T)` datasets) to `get_video_file_reader`; reads are permuted back to `(T, H, W, C)`. ### TIFF (.tif, .tiff) diff --git a/docs/user_guide/online_processing.md b/docs/user_guide/online_processing.md index 908c007..8eac91e 100644 --- a/docs/user_guide/online_processing.md +++ b/docs/user_guide/online_processing.md @@ -41,11 +41,11 @@ flow_reg.set_reference() Behavior by input shape: -- **4D stack `(T, H, W, C)`** -- the frames are preregistered with `compensate_arr` against their temporal mean, using a copy of the options with `quality_setting="balanced"`. The mean of the registered frames becomes the reference. -- **3D `(H, W, C)` or 2D `(H, W)`** -- used directly as a single-frame reference, without preregistration. A 3D input is always interpreted as one multi-channel frame, so a grayscale stack `(T, H, W)` must be given an explicit channel axis first (e.g. `frames[..., None]`). +- **Stack `(T, H, W, C)` or `(T, H, W)`** -- the frames are preregistered with `compensate_arr` against their temporal mean, using a copy of the options with `quality_setting="balanced"`. The mean of the registered frames becomes the reference. +- **Single frame `(H, W, C)` or `(H, W)`** -- used directly as the reference, without preregistration. A 3D input is interpreted as one multi-channel frame when its last dimension is at most 4 (the `ArrayReader` convention), and as a grayscale `(T, H, W)` stack otherwise. - **No argument** -- the frames accumulated in the internal reference buffer are stacked and processed as above. Raises `ValueError` if the buffer is empty. -After the reference is established, it is normalized and spatially Gaussian-filtered; the min/max of this filtered reference is stored and used to normalize every incoming frame (per channel when `channel_normalization="separate"`, jointly otherwise). Setting a reference also resets the stored flow initialization and clears the temporal filter buffer. +After the reference is established, it is normalized and spatially Gaussian-filtered; the min/max of the raw reference is stored and used to normalize every incoming frame (per channel when `channel_normalization="separate"`, jointly otherwise) -- the same convention as the batch pipeline, so frames can arrive at the recording's native scale (e.g. uint16). Setting a reference also resets the stored flow initialization and clears the temporal filter buffer. `reset_reference(new_reference)` is an alias that calls `set_reference()` with the given array. @@ -65,14 +65,12 @@ If no reference has been set yet, the call instead appends the frame to the refe Once a reference exists, each call: -1. Normalizes the frame using the stored reference min/max values. +1. Normalizes the frame using the raw reference's min/max values. 2. Applies the 2D spatial Gaussian filter and pushes the result into the temporal buffer. 3. Applies the causal temporal half-kernel filter (see below). 4. Computes the displacement field against the filtered reference, initialized with the previous frame's flow. 5. Warps the original (unfiltered) input frame with the resulting field, using `options.interpolation_method` (default cubic). -6. Every `reference_update_interval`-th frame, blends the warped frame into the reference with weight `reference_update_weight` and refreshes the normalization min/max from the updated reference. - -The `normalize` keyword of `__call__` is currently unused; normalization is always applied. +6. Every `reference_update_interval`-th frame, blends the warped frame into the reference with weight `reference_update_weight` and refreshes the normalization min/max from the updated raw reference. Convenience methods: @@ -88,7 +86,7 @@ Batch processing filters each batch with a 3D Gaussian over `(y, x, t)`. In a st 2. The 2D-filtered frames are kept in a circular buffer of size `max(1, int(truncate * sigma_t + 0.5) + 1)`, where `sigma_t` is the temporal sigma. 3. The frame used for flow estimation is a weighted average of the current and past buffered frames with normalized half-Gaussian weights (`pyflowreg.util.image_processing.gaussian_filter_1d_half_kernel`). Only the current and past frames contribute, so the filter is causal and adds no latency, but it is not identical to the symmetric temporal filtering used in batch mode. -With the `OFOptions` default `sigma` of `[[1.0, 1.0, 0.1], [1.0, 1.0, 0.1]]`, the temporal sigma of 0.1 yields a buffer of size 1 and no effective temporal filtering; increase the third sigma component to enable it. When per-channel sigmas are configured, `FlowRegLive` currently uses the first channel's spatial sigmas for all channels and the maximum temporal sigma across channels. +With the `OFOptions` default `sigma` of `[[1.0, 1.0, 0.1], [1.0, 1.0, 0.1]]`, the temporal sigma of 0.1 yields a buffer of size 1 and no effective temporal filtering; increase the third sigma component to enable it. When per-channel sigmas are configured, each channel is spatially filtered with its own sigmas; the temporal half-kernel uses the maximum temporal sigma across channels. ## Complete Example diff --git a/docs/user_guide/prealignment.md b/docs/user_guide/prealignment.md index e0424e6..df49bdf 100644 --- a/docs/user_guide/prealignment.md +++ b/docs/user_guide/prealignment.md @@ -71,12 +71,14 @@ the result. `compensate_recording` and `compensate_arr`. All three executors (sequential, threading, multiprocessing) implement it; see [Parallelization](parallelization.md) for executor selection. +- Reference preregistration in `OFOptions.get_reference_frame` forwards the + `cc_*` settings, so an enabled pre-alignment also applies while building + the reference. (The MATLAB reference does not pre-align its + preregistration; this is a deliberate improvement.) - It is **not** applied by: - Direct `get_displacement` calls — the flow backends do not accept the `cc_*` parameters; the executors consume them before invoking the solver. - `FlowRegLive` online processing, which has no pre-alignment handling. - - Reference preregistration in `OFOptions.get_reference_frame`, which builds - its own internal options without `cc_initialization`. - Each pre-aligned frame requires two additional warps and one FFT-based correlation, so leave it disabled when frame-to-frame motion is small. diff --git a/docs/user_guide/workflows.md b/docs/user_guide/workflows.md index 37366b3..99bd2d6 100644 --- a/docs/user_guide/workflows.md +++ b/docs/user_guide/workflows.md @@ -162,7 +162,7 @@ processor.set_reference(corrected_frames) processor.reset_reference(new_reference_frame) ``` -Note that `set_reference()` interprets a 3D array as a single multi-channel frame `(H, W, C)`, not as a grayscale stack. Calling `set_reference()` without arguments uses the frames buffered internally before a reference existed. +Note that `set_reference()` interprets a 3D array as a single multi-channel frame `(H, W, C)` when its last dimension is at most 4 (the `ArrayReader` convention), and as a grayscale `(T, H, W)` stack otherwise. Calling `set_reference()` without arguments uses the frames buffered internally before a reference existed. ### When to Use diff --git a/requirements_win.txt b/requirements_win.txt index 212b2a8..03d3f29 100644 --- a/requirements_win.txt +++ b/requirements_win.txt @@ -8,3 +8,6 @@ numba pywin32 pydantic hdf5storage +pyyaml +Pillow +tomli diff --git a/src/pyflowreg/core/__init__.py b/src/pyflowreg/core/__init__.py index 4c456ee..3eb713d 100644 --- a/src/pyflowreg/core/__init__.py +++ b/src/pyflowreg/core/__init__.py @@ -35,6 +35,8 @@ List all available backends is_backend_available Check if a specific backend is available +get_backend_executors + Get the parallelization executors supported by a backend See Also -------- @@ -54,6 +56,7 @@ get_backend, list_backends, is_backend_available, + get_backend_executors, ) __all__ = [ @@ -62,6 +65,7 @@ "get_backend", "list_backends", "is_backend_available", + "get_backend_executors", ] # Register built-in backends diff --git a/src/pyflowreg/core/diso_optical_flow.py b/src/pyflowreg/core/diso_optical_flow.py index 295f74e..29b5c30 100644 --- a/src/pyflowreg/core/diso_optical_flow.py +++ b/src/pyflowreg/core/diso_optical_flow.py @@ -30,9 +30,11 @@ ``gradient_descent_iterations``, ``patch_size``, ``patch_stride``, ``use_mean_normalization``, ``use_spatial_propagation``). - Variational solver keywords forwarded by the pipelines (``alpha``, - ``iterations``, ``uv``, ``const_assumption``, ``gnc_schedule``, ...) are + ``iterations``, ``const_assumption``, ``gnc_schedule``, ...) are accepted by :meth:`DisoOF.__call__` for signature compatibility but - ignored. + ignored. The flow initialization keyword ``uv`` is the exception: it is + honored as the DIS warm start (``w`` takes precedence when both are + given). References ---------- @@ -339,16 +341,21 @@ def __call__( conventions as ``fixed``. w : ndarray, optional Initial displacement field of shape (H, W, 2) used as warm - start. Used only if it is an ndarray with exactly this layout; - it is cast to float32. + start. Takes precedence over ``uv`` when both are given. Used + only if it is an ndarray with exactly this layout and the same + spatial size as the images; it is converted to contiguous + float32. weight : ndarray, optional Channel weights for multi-channel inputs; see ``_to_gray`` for the accepted formats. **kwargs : dict Accepted for signature compatibility with the variational - :func:`pyflowreg.core.optical_flow.get_displacement` (e.g. - ``alpha``, ``iterations``, ``uv``, ``const_assumption``, - ``gnc_schedule``); all extra keyword arguments are ignored. + :func:`pyflowreg.core.optical_flow.get_displacement`. The + ``uv`` keyword (the flow initialization passed by the batch + pipelines) is honored as warm start when ``w`` is not given; + all other extra keyword arguments (e.g. ``alpha``, + ``iterations``, ``const_assumption``, ``gnc_schedule``) are + ignored. Returns ------- @@ -362,8 +369,9 @@ def __call__( Notes ----- The batch pipelines pass the flow initialization as the keyword - ``uv``, which this wrapper does not read; a warm start is only used - when supplied via ``w``. + ``uv``; it is forwarded to ``cv2.DISOpticalFlow.calc`` as the + initial flow. A warm start whose spatial size does not match the + images is silently skipped (mirroring OpenCV's own size check). """ self._ensure() @@ -374,15 +382,21 @@ def __call__( # Normalize to [0,255] uint8 A, B = self._normalize(a, b) - # Prepare initial flow if provided + # Prepare initial flow if provided. Precedence: explicit ``w`` wins, + # otherwise fall back to the pipeline keyword ``uv`` (used by the + # executors, FlowRegLive and compensate_recording). init = None + init_src = w if w is not None else kwargs.get("uv") if ( - w is not None - and isinstance(w, np.ndarray) - and w.ndim == 3 - and w.shape[2] == 2 + isinstance(init_src, np.ndarray) + and init_src.ndim == 3 + and init_src.shape[2] == 2 + and init_src.shape[:2] == A.shape[:2] ): - init = w.astype(np.float32, copy=False) + # Copy: cv2.DISOpticalFlow.calc treats the init flow as an + # InputOutputArray and writes the result into it in place; cv2 + # also requires contiguous float32 (CV_32FC2) of image size. + init = np.array(init_src, dtype=np.float32, order="C", copy=True) # Compute optical flow flow = self._dis.calc(A, B, init) diff --git a/src/pyflowreg/motion_correction/OF_options.py b/src/pyflowreg/motion_correction/OF_options.py index 1cdb088..4019462 100644 --- a/src/pyflowreg/motion_correction/OF_options.py +++ b/src/pyflowreg/motion_correction/OF_options.py @@ -325,7 +325,10 @@ class OFOptions(BaseModel): - ``reference_frames`` (list of int, str, Path or ndarray, default ``list(range(50, 500))``): Frame indices to preregister and average, - an image file path, or a precomputed reference array. + an image file path, or a precomputed reference array. Indices beyond + the recording length are clipped to the last frame and the resulting + duplicates are removed with a printed warning, so short recordings + preregister each available frame exactly once. - ``update_reference`` (bool, default ``False``): Update reference during processing. - ``n_references`` (int, default ``1``): Number of references; @@ -337,7 +340,9 @@ class OFOptions(BaseModel): **Pre-alignment** - ``cc_initialization`` (bool, default ``False``): Enable - cross-correlation initialization. + cross-correlation initialization. Also applied during reference + preregistration in ``get_reference_frame`` (an improvement over the + MATLAB reference, whose preregistration is not pre-aligned). - ``cc_hw`` (int or 2-tuple of int, default ``256``): Target height/width for cross-correlation projections. - ``cc_up`` (int, default ``1``): Upsampling factor for subpixel @@ -790,7 +795,12 @@ def get_reference_frame( video_reader: Optional[VideoReader] = None, registration_config: Optional[Any] = None, ) -> Union[np.ndarray, List[np.ndarray]]: - """Get reference frame(s), with optional preregistration.""" + """Get reference frame(s), with optional preregistration. + + Index lists are clipped to the recording length (duplicates from + clipping are removed with a printed warning); the preregistration + honors ``cc_initialization``/``cc_hw``/``cc_up``. + """ if self.n_references > 1: warnings.warn( "Multi-reference mode not fully implemented; repeating a single computed reference" @@ -833,10 +843,20 @@ def get_reference_frame( valid_indices.append(idx) if clipped: + # Deduplicate (order-preserving): clipping maps every + # out-of-range index to the last frame, which would + # otherwise be read, flow-solved, and averaged hundreds of + # times, massively over-weighting it in the reference. + # Intentional duplicates in fully in-range user lists are + # preserved (no dedup when nothing was clipped). + n_requested = len(valid_indices) + valid_indices = list(dict.fromkeys(valid_indices)) + n_removed = n_requested - len(valid_indices) print( f"Warning: Reference frames exceed video length ({frame_count} frames). " f"Clipping indices from {self.reference_frames[0]}-{self.reference_frames[-1]} " - f"to {valid_indices[0]}-{valid_indices[-1]}" + f"to {valid_indices[0]}-{valid_indices[-1]} and removing " + f"{n_removed} duplicate indices." ) frames = video_reader[valid_indices] # (T,H,W,C) using array-like indexing @@ -867,9 +887,13 @@ def get_reference_frame( if gaussian_filter is not None: frames_smooth = np.zeros_like(frames) for c in range(n_channels): + # sig is (sx, sy, st); the per-channel slice is + # (H, W, T), so reorder to scipy's axis order (sy, sx, st). sig = self.get_sigma_at(c) + np.array([1, 1, 0.5]) frames_smooth[:, :, c, :] = gaussian_filter( - frames[:, :, c, :], sigma=tuple(sig), mode="reflect" + frames[:, :, c, :], + sigma=(sig[1], sig[0], sig[2]), + mode="reflect", ) else: frames_smooth = frames @@ -900,7 +924,11 @@ def get_reference_frame( else self.alpha + 2.0 ) - # Create a temporary OFOptions for preregistration + # Create a temporary OFOptions for preregistration. The + # cross-correlation pre-alignment settings are forwarded so an + # enabled cc_initialization also applies during reference + # preregistration (the MATLAB reference does not pre-align its + # preregistration; this is a deliberate improvement). prereg_options = OFOptions( alpha=alpha_prereg, levels=self.levels, @@ -913,6 +941,9 @@ def get_reference_frame( constancy_assumption=self.constancy_assumption, weight=weight_2d, buffer_size=self.buffer_size, + cc_initialization=self.cc_initialization, + cc_hw=self.cc_hw, + cc_up=self.cc_up, ) # Reshape frames_norm from (H,W,C,T) to (T,H,W,C) for compensate_arr diff --git a/src/pyflowreg/motion_correction/compensate_arr.py b/src/pyflowreg/motion_correction/compensate_arr.py index 7ab837f..8e0ea46 100644 --- a/src/pyflowreg/motion_correction/compensate_arr.py +++ b/src/pyflowreg/motion_correction/compensate_arr.py @@ -41,9 +41,9 @@ def compensate_arr( c1 : ndarray Input array to register, shape (T, H, W, C), (T, H, W), (H, W, C), or (H, W). A 3D input is treated as (T, H, W) when ``c_ref`` is 2D; - otherwise the array reader treats a 3D input as a single (H, W, C) - frame when the last dimension is at most 4, and as (T, H, W) - when it is larger. + otherwise it is treated as a single (H, W, C) frame when the last + dimension is at most 4 (the ArrayReader convention), and as + (T, H, W) when it is larger. c_ref : ndarray Reference frame, shape (H, W, C) or (H, W). options : OFOptions, optional @@ -82,18 +82,18 @@ def compensate_arr( Returns ------- c_reg : ndarray - Registered frames. Inputs of shape (T, H, W, C), (T, H, W), and - (H, W) keep their input shape; an (H, W, C) input is returned as - (1, H, W, C). Cast according to ``options.output_typename`` - (default "double", i.e. float64) when it is one of "single", - "double", "uint8", "uint16", "int16", "int32"; other values leave - the pipeline output dtype unchanged. + Registered frames, with the same shape (rank) as the input: + (T, H, W, C), (T, H, W), (H, W, C), or (H, W). Cast according to + ``options.output_typename`` (default "double", i.e. float64) when + it is one of "single", "double", "uint8", "uint16", "int16", + "int32"; other values leave the pipeline output dtype unchanged. w : ndarray Displacement fields, shape (T, H, W, 2) with components (u, v), where ``w[..., 0]`` is the horizontal (x) and ``w[..., 1]`` the - vertical (y) displacement. For a single 2D (H, W) input the leading - time axis is removed, giving (H, W, 2). A zero-filled array is - returned if no displacement fields were captured. + vertical (y) displacement. For single-frame inputs ((H, W) or + (H, W, C)) the leading time axis is removed, giving (H, W, 2). A + zero-filled array is returned if no displacement fields were + captured. Raises ------ @@ -121,7 +121,8 @@ def compensate_arr( >>> registered, flow = compensate_arr(video, reference, progress_callback=progress) # doctest: +SKIP """ # Handle 3D squeeze for single channel (MATLAB compatibility) - squeezed = False + squeezed = False # channel axis was added + single_frame = False # time axis was added (input had no T axis) original_shape = c1.shape # Validate input is not empty @@ -133,12 +134,20 @@ def compensate_arr( c1 = c1[..., np.newaxis] c_ref = c_ref[..., np.newaxis] squeezed = True + elif c1.ndim == 3 and c1.shape[-1] <= 4: + # Single (H, W, C) frame with a multi-channel reference - add the + # time axis here (mirrors the ArrayReader heuristic: a 3D array + # with last dimension <= 4 is one multi-channel frame) so the + # output rank can mirror the input rank. + c1 = c1[np.newaxis, ...] + single_frame = True elif c1.ndim == 2: # Single frame, single channel c1 = c1[np.newaxis, :, :, np.newaxis] if c_ref.ndim == 2: c_ref = c_ref[..., np.newaxis] squeezed = True + single_frame = True # Configure options for array processing if options is None: @@ -154,7 +163,9 @@ def compensate_arr( if options.backend_params: options.backend_params.update(backend_params) else: - options.backend_params = backend_params + # Defensive copy so the internal options never alias the + # caller's kwarg dict. + options.backend_params = dict(backend_params) if get_displacement is not None: options.get_displacement_impl = get_displacement if get_displacement_factory is not None: @@ -208,25 +219,32 @@ def compensate_arr( if options.output_typename in dtype_map: c_reg = c_reg.astype(dtype_map[options.output_typename]) - # Squeeze back if needed to match input shape + # Squeeze back so the output rank mirrors the input rank if squeezed: if len(original_shape) == 2: - # Was single frame (H,W) - c_reg = np.squeeze(c_reg) - if w is not None: - w = np.squeeze(w, axis=0) # Remove time dimension + # Was single frame (H,W): drop time and channel axes + c_reg = np.squeeze(c_reg, axis=(0, -1)) elif len(original_shape) == 3: - # Was (T,H,W) or (H,W,C) - c_reg = np.squeeze(c_reg, axis=-1) # Remove channel dimension - - # If no flow fields were captured, create empty array + # Was (T,H,W) with a 2D reference: drop the channel axis + c_reg = np.squeeze(c_reg, axis=-1) + elif single_frame: + # Was a single (H,W,C) frame: drop the time axis + c_reg = np.squeeze(c_reg, axis=0) + + if w is not None and single_frame: + w = np.squeeze(w, axis=0) # Remove time dimension + + # If no flow fields were captured, create a zero-filled array matching + # the documented shape: (H, W, 2) for single-frame inputs, (T, H, W, 2) + # otherwise if w is None: - if c_reg.ndim >= 3: - T = c_reg.shape[0] if c_reg.ndim == 4 else 1 - H, W = c_reg.shape[-3:-1] if c_reg.ndim == 4 else c_reg.shape[:2] + if single_frame: + H, W = original_shape[:2] + w = np.zeros((H, W, 2), dtype=np.float32) else: - T, H, W = 1, c_reg.shape[0], c_reg.shape[1] - w = np.zeros((T, H, W, 2), dtype=np.float32) + # (T, H, W, C) or (T, H, W): spatial size at axes 1-2 + H, W = original_shape[1:3] + w = np.zeros((original_shape[0], H, W, 2), dtype=np.float32) return c_reg, w diff --git a/src/pyflowreg/motion_correction/compensate_recording.py b/src/pyflowreg/motion_correction/compensate_recording.py index 9d431ac..8fc2198 100644 --- a/src/pyflowreg/motion_correction/compensate_recording.py +++ b/src/pyflowreg/motion_correction/compensate_recording.py @@ -69,10 +69,12 @@ class RegistrationConfig: back to 4 if the core count cannot be determined). Any other value is used directly as the worker count. verbose : bool, optional - Verbosity flag, default ``False``. Note that in the current - implementation most pipeline progress messages (executor - selection, batch timings, final statistics) are printed when - this is ``False``. + Verbosity flag, default ``False``. When ``True``, pipeline + progress messages (executor selection, batch timings, final + statistics) are printed; warnings are always emitted regardless + of this flag. Note: the MATLAB reference gates the same messages + on ``~options.verbose`` (chatty by default); PyFlowReg + deliberately uses the natural polarity and is quiet by default. parallelization : str or None, optional Name of the parallelization executor: ``'sequential'``, ``'threading'``, or ``'multiprocessing'``. ``None`` (default) @@ -290,10 +292,9 @@ def _setup_executor(self): executor_class = RuntimeContext.get_parallelization_executor(executor_name) if executor_class is None: # Fallback to sequential if requested executor not available - if not self.config.verbose: - print( - f"Warning: {executor_name} executor not available, falling back to sequential" - ) + warnings.warn( + f"{executor_name} executor not available, falling back to sequential" + ) executor_class = RuntimeContext.get_parallelization_executor("sequential") # If sequential is also not available, import and register it @@ -316,7 +317,7 @@ def _setup_executor(self): # Create executor instance self.executor = executor_class(n_workers=self.n_workers) - if not self.config.verbose: + if self.config.verbose: # Use actual executor name and worker count actual_workers = self.executor.n_workers worker_str = "worker" if actual_workers == 1 else "workers" @@ -648,7 +649,7 @@ def _compute_initial_w( """Compute initial displacement field from first frames.""" n_init = min(22, first_batch.shape[0]) # T is first dimension - if not self.config.verbose: + if self.config.verbose: print("Computing initial displacement field...") # Process first n_init frames - use "initial_w" task to avoid counting toward main progress @@ -672,7 +673,7 @@ def _compute_initial_w( # Average flows w_init = np.mean(w, axis=0) - if not self.config.verbose: + if self.config.verbose: print("Done pre-registration to get w_init.") return w_init @@ -757,7 +758,7 @@ def run(self, reference_frame: Optional[np.ndarray] = None) -> np.ndarray: # Initialize total frames for progress tracking self._total_frames = len(self.video_reader) if self.video_reader else None - if not self.config.verbose: + if self.config.verbose: quality = getattr(self.options, "quality_setting", "balanced") print(f"\nStarting compensation with quality={quality}") print( @@ -870,7 +871,7 @@ def run(self, reference_frame: Optional[np.ndarray] = None) -> np.ndarray: total_frames += registered.shape[0] batch_time = time() - batch_start - if not self.config.verbose: + if self.config.verbose: fps = registered.shape[0] / batch_time print( f"Batch {batch_idx}: {registered.shape[0]} frames in {batch_time:.2f}s ({fps:.1f} fps)" @@ -883,7 +884,7 @@ def run(self, reference_frame: Optional[np.ndarray] = None) -> np.ndarray: # Final stats total_time = time() - start_time - if not self.config.verbose: + if self.config.verbose: avg_fps = total_frames / max(1e-6, total_time) print( f"\nProcessed {total_frames} frames in {total_time:.2f}s (avg {avg_fps:.1f} fps)" @@ -919,7 +920,8 @@ def _save_metadata(self): ref_path = output_path / "reference_frame.npy" np.save(str(ref_path), self.reference_raw) - print(f"Saved metadata to {output_path}") + if self.config.verbose: + print(f"Saved metadata to {output_path}") def _cleanup(self): """Close file handlers.""" diff --git a/src/pyflowreg/motion_correction/flow_reg_live.py b/src/pyflowreg/motion_correction/flow_reg_live.py index c7c369f..5c53d8c 100644 --- a/src/pyflowreg/motion_correction/flow_reg_live.py +++ b/src/pyflowreg/motion_correction/flow_reg_live.py @@ -77,14 +77,26 @@ def __init__( # Get sigma values and determine temporal buffer size self.sigma = np.asarray(self.options.sigma) if self.sigma.ndim == 1: - self.sigma_2d = self.sigma[:2] # [sy, sx] + self.sigma_2d = self.sigma[:2] # [sx, sy] self.sigma_t = self.sigma[2] if len(self.sigma) > 2 else 0.0 else: - # Per-channel, find max temporal sigma - self.sigma_2d = self.sigma[0, :2] # Use first channel for 2D - # Get max temporal sigma across all channels + # Per-channel spatial sigmas (C, 2) — apply_gaussian_filter + # consumes per-channel rows natively, so every channel keeps + # its own [sx, sy]. + self.sigma_2d = self.sigma[:, :2] + # The temporal half-kernel filter takes a scalar sigma; use the + # max across channels (also sizes the temporal buffer). self.sigma_t = np.max(self.sigma[:, 2]) if self.sigma.shape[1] > 2 else 0.0 + # channel_normalization mode for util.image_processing.normalize, + # which expects "together"/"separate" (OFOptions uses "joint" for + # the joint mode). + self._norm_mode = ( + "separate" + if self.options.channel_normalization.value == "separate" + else "together" + ) + # Calculate temporal buffer size from sigma_t # Buffer size = kernel radius + 1, where radius = truncate * sigma_t temporal_buffer_size = max(1, int(self.truncate * self.sigma_t + 0.5) + 1) @@ -104,7 +116,8 @@ def __init__( # Flow initialization self.last_flow = None - # Normalization parameters from filtered reference + # Normalization parameters from the RAW reference (same bounds the + # batch pipeline uses to normalize incoming raw frames) self.norm_min = None self.norm_max = None @@ -113,18 +126,25 @@ def set_reference(self, frames: Optional[np.ndarray] = None): Initialize reference from frames or buffer. Args: - frames: Optional array of frames (T,H,W,C). If None, uses buffer. + frames: Optional reference input. A (T,H,W,C) or (T,H,W) stack + is compensated with compensate_arr and averaged; a single + (H,W,C) or (H,W) frame is used directly. A 3D array is + interpreted as a single (H,W,C) frame when its last + dimension is at most 4 (the ArrayReader convention), + otherwise as a (T,H,W) stack. If None, uses the buffer. """ if frames is not None: # Use provided frames if frames.ndim == 2: # Single 2D frame (H,W) - convert to 3D self.reference_raw = frames[..., np.newaxis].copy() - elif frames.ndim == 3: - # Single frame (H,W,C) + elif frames.ndim == 3 and frames.shape[-1] <= 4: + # Single frame (H,W,C) — ArrayReader convention: a 3D array + # with last dimension <= 4 is one multi-channel frame, a + # larger last dimension means a grayscale (T,H,W) stack. self.reference_raw = frames.copy() else: - # Multiple frames - compensate and average + # Stack (T,H,W,C) or (T,H,W) - compensate and average print(f"Preregistering {frames.shape[0]} frames for reference...") # Use balanced quality for reference @@ -159,16 +179,8 @@ def set_reference(self, frames: Optional[np.ndarray] = None): # Preprocess reference (normalize and filter) self._preprocess_reference() - # Store normalization parameters from FILTERED reference - if self.options.channel_normalization.value == "separate": - self.norm_min = [] - self.norm_max = [] - for c in range(self.reference_proc.shape[2]): - self.norm_min.append(self.reference_proc[..., c].min()) - self.norm_max.append(self.reference_proc[..., c].max()) - else: - self.norm_min = self.reference_proc.min() - self.norm_max = self.reference_proc.max() + # Store normalization parameters from the RAW reference + self._update_norm_bounds() # Reset flow initialization self.last_flow = None @@ -176,12 +188,31 @@ def set_reference(self, frames: Optional[np.ndarray] = None): # Clear temporal buffer self.temporal_buffer.clear() + def _update_norm_bounds(self): + """Store normalization bounds from the RAW reference. + + The batch pipeline normalizes raw incoming frames with the raw + reference's range (``normalization_ref=reference_raw`` in + ``compensate_recording``); these are the same bounds the reference + itself is normalized with in ``_preprocess_reference``, so frames + and reference reach the flow solver on the same scale. + """ + if self._norm_mode == "separate": + self.norm_min = [] + self.norm_max = [] + for c in range(self.reference_raw.shape[2]): + self.norm_min.append(self.reference_raw[..., c].min()) + self.norm_max.append(self.reference_raw[..., c].max()) + else: + self.norm_min = self.reference_raw.min() + self.norm_max = self.reference_raw.max() + def _preprocess_reference(self): """Preprocess reference frame (normalize and filter).""" # First normalize ref_norm = normalize( self.reference_raw, - channel_normalization=self.options.channel_normalization.value, + channel_normalization=self._norm_mode, ) # Apply 2D Gaussian filter @@ -190,15 +221,13 @@ def _preprocess_reference(self): ) def _normalize_with_reference(self, frame: np.ndarray) -> np.ndarray: - """Normalize frame using filtered reference min/max values.""" + """Normalize frame using the raw reference min/max values.""" if self.norm_min is None or self.norm_max is None: # Fallback to frame's own normalization - return normalize( - frame, channel_normalization=self.options.channel_normalization.value - ) + return normalize(frame, channel_normalization=self._norm_mode) eps = 1e-8 - if self.options.channel_normalization.value == "separate": + if self._norm_mode == "separate": result = np.zeros_like(frame, dtype=np.float64) for c in range(frame.shape[-1]): if c < len(self.norm_min): @@ -213,15 +242,16 @@ def _normalize_with_reference(self, frame: np.ndarray) -> np.ndarray: else: return (frame - self.norm_min) / (self.norm_max - self.norm_min + eps) - def __call__( - self, frame: np.ndarray, normalize: bool = True - ) -> Tuple[np.ndarray, np.ndarray]: + def __call__(self, frame: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: """ Process single frame with motion compensation. + The frame is normalized against the raw reference's intensity + range, so it can arrive at the recording's native scale (e.g. + uint16). + Args: frame: Input frame (H,W,C) or (H,W) - normalize: Whether to normalize (unused, always normalizes) Returns: Tuple of (registered_frame, flow_field) with shapes (H,W,C) and (H,W,2) @@ -237,7 +267,7 @@ def __call__( self.frame_count += 1 - # Normalize using filtered reference parameters + # Normalize using the raw reference's range frame_norm = self._normalize_with_reference(frame) # Apply 2D Gaussian filter @@ -309,16 +339,8 @@ def __call__( 1 - self.reference_update_weight ) * self.reference_raw + self.reference_update_weight * registered - # Update normalization parameters from new filtered reference - if self.options.channel_normalization.value == "separate": - self.norm_min = [] - self.norm_max = [] - for c in range(self.reference_proc.shape[2]): - self.norm_min.append(self.reference_proc[..., c].min()) - self.norm_max.append(self.reference_proc[..., c].max()) - else: - self.norm_min = self.reference_proc.min() - self.norm_max = self.reference_proc.max() + # Update normalization parameters from the new raw reference + self._update_norm_bounds() return registered, flow @@ -336,7 +358,7 @@ def register_frames(self, frames: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: flows = np.empty(frames.shape[:3] + (2,), dtype=np.float32) for i in range(frames.shape[0]): - registered[i], flows[i] = self(frames[i], False) + registered[i], flows[i] = self(frames[i]) return registered, flows diff --git a/src/pyflowreg/session/stage1_compensate.py b/src/pyflowreg/session/stage1_compensate.py index c030a35..a0c8bf6 100644 --- a/src/pyflowreg/session/stage1_compensate.py +++ b/src/pyflowreg/session/stage1_compensate.py @@ -417,7 +417,9 @@ def compensate_single_recording( compensated_h5 = next((p for p in candidates if p.exists()), candidates[0]) if not compensated_h5.exists(): - reg_config = RegistrationConfig(n_jobs=config.n_workers) + # verbose=True keeps the per-batch console progress for long + # session runs (the pipeline is quiet by default). + reg_config = RegistrationConfig(n_jobs=config.n_workers, verbose=True) compensate_recording(options, config=reg_config) # Re-check both candidates after compensation compensated_h5 = next((p for p in candidates if p.exists()), candidates[0]) diff --git a/src/pyflowreg/util/image_processing.py b/src/pyflowreg/util/image_processing.py index 054f52d..fa46efd 100644 --- a/src/pyflowreg/util/image_processing.py +++ b/src/pyflowreg/util/image_processing.py @@ -114,9 +114,13 @@ def apply_gaussian_filter( Input array, shape (H, W, C) or (T, H, W, C). sigma : np.ndarray Standard deviations of the Gaussian kernel ordered as - ``[sy, sx, st]``. Shape (3,) applies the same sigmas to all - channels; shape (n_channels, 3) gives per-channel sigmas. For - (H, W, C) input only the spatial components ``[sy, sx]`` are used. + ``[sx, sy, st]`` — ``sx`` smooths along the width/columns (W), + ``sy`` along the height/rows (H), ``st`` along frames (T) — + matching MATLAB's ``imgaussfilt``/``imgaussfilt3`` convention and + the ``OFOptions.sigma`` field. Shape (3,) applies the same sigmas + to all channels; shape (n_channels, 3) gives per-channel sigmas. + For (H, W, C) input only the spatial components ``[sx, sy]`` are + used. mode : str, optional Boundary handling mode passed to ``scipy.ndimage.gaussian_filter``. Default is "reflect". @@ -139,8 +143,9 @@ def apply_gaussian_filter( s = sigma[min(c, len(sigma) - 1), :2] # Use only spatial components else: s = sigma[:2] # Use first two components + # Reorder from (sx, sy) to scipy's axis order (sy, sx) for (H, W) result[..., c] = gaussian_filter( - arr[..., c], sigma=s, mode=mode, truncate=truncate + arr[..., c], sigma=(s[1], s[0]), mode=mode, truncate=truncate ) return result @@ -149,10 +154,11 @@ def apply_gaussian_filter( for c in range(arr.shape[3]): # C is last dimension if sigma.ndim == 2: # Per-channel sigmas s = sigma[min(c, len(sigma) - 1)] - # Reorder from (sy, sx, st) to (st, sy, sx) for scipy - s_3d = (s[2], s[0], s[1]) else: - s_3d = (sigma[2], sigma[0], sigma[1]) + s = sigma + # Reorder from (sx, sy, st) to scipy's axis order (st, sy, sx) + # for the (T, H, W) per-channel volume + s_3d = (s[2], s[1], s[0]) # Apply 3D Gaussian filter result[..., c] = gaussian_filter( diff --git a/src/pyflowreg/util/io/_base.py b/src/pyflowreg/util/io/_base.py index a637ced..28d1163 100644 --- a/src/pyflowreg/util/io/_base.py +++ b/src/pyflowreg/util/io/_base.py @@ -316,6 +316,10 @@ def has_batch(self) -> bool: """ Check whether more frames are available for sequential reading. + Initializes the reader on first use, so a freshly constructed + reader can be consumed directly with ``while reader.has_batch():`` + or ``for batch in reader:``. + Returns ------- bool @@ -323,6 +327,7 @@ def has_batch(self) -> bool: frames have been consumed by ``read_batch()``; False once the end of the file is reached. """ + self._ensure_initialized() return self.current_frame < self.frame_count def reset(self): diff --git a/src/pyflowreg/util/io/hdf5.py b/src/pyflowreg/util/io/hdf5.py index df2d699..4c48ea8 100644 --- a/src/pyflowreg/util/io/hdf5.py +++ b/src/pyflowreg/util/io/hdf5.py @@ -15,7 +15,8 @@ class HDF5FileReader(DSFileReader, VideoReader): Reads one or more HDF5 datasets as video channels. If no ``dataset_names`` are given, suitable datasets are discovered with the DSFileReader heuristic. 3D datasets are interpreted as (T, H, W) with - one dataset per channel; a 4D dataset is interpreted as (T, H, W, C). + one dataset per channel by default (configurable via + ``dimension_ordering``); a 4D dataset is interpreted as (T, H, W, C). Frames are returned in (T, H, W, C) format. """ @@ -39,6 +40,13 @@ def __init__( - ``dataset_names`` (list of str): Datasets to read, one per channel. If not given, datasets are discovered automatically. + - ``dimension_ordering`` (tuple of int): Axis positions of + (height, width, time) in each stored 3D dataset, matching + ``HDF5FileWriter`` and ``MATFileReader``. Default ``(1, 2, 0)`` + means the dataset is stored as (T, H, W). Ignored for 4D + datasets, which are always read as (T, H, W, C). Note that + indexing a non-leading time axis reads across HDF5 chunks + and is slower than the default layout. """ # Initialize parent classes DSFileReader.__init__(self) @@ -51,7 +59,12 @@ def __init__( # Dataset-specific options self.dataset_names = kwargs.get("dataset_names") - self.dimension_ordering = kwargs.get("dimension_ordering") + self.dimension_ordering = tuple(kwargs.get("dimension_ordering", (1, 2, 0))) + if sorted(self.dimension_ordering) != [0, 1, 2]: + raise ValueError( + "dimension_ordering must be a permutation of (0, 1, 2) giving " + f"the axis positions of (H, W, T), got {self.dimension_ordering}" + ) def _initialize(self): """Open file and set up properties.""" @@ -78,12 +91,16 @@ def visitor(name, obj): first_ds = self.h5file[self.dataset_names[0]] shape = first_ds.shape - # Detect dimension ordering (implementation specific) - # For now assume it's already (T, H, W) or needs transposing + # Interpret the stored layout via dimension_ordering: positions of + # (H, W, T) in the dataset shape; default (1, 2, 0) = (T, H, W). if len(shape) == 3: - self.frame_count, self.height, self.width = shape + self.height = shape[self.dimension_ordering[0]] + self.width = shape[self.dimension_ordering[1]] + self.frame_count = shape[self.dimension_ordering[2]] self.n_channels = len(self.dataset_names) elif len(shape) == 4: + # 4D datasets are always (T, H, W, C); dimension_ordering only + # applies to per-channel 3D datasets. self.frame_count, self.height, self.width, self.n_channels = shape self.dtype = first_ds.dtype @@ -127,16 +144,33 @@ def _read_raw_frames(self, frame_indices: Union[slice, List[int]]) -> np.ndarray (n_frames, self.height, self.width, self.n_channels), dtype=self.dtype ) - # Read from each dataset/channel + # Read from each dataset/channel, indexing the configured time axis + # and permuting the stored (H, W, T) positions to (T, H, W) — + # mirrors MATFileReader. + do = self.dimension_ordering + t_axis = do[2] for ch_idx, ds_name in enumerate(self.dataset_names): dataset = self.h5file[ds_name] - if isinstance(frame_indices, slice): - # Efficient slicing for contiguous frames - data = dataset[frame_indices, :, :] + if dataset.ndim == 3: + idx = [slice(None), slice(None), slice(None)] + if isinstance(frame_indices, slice): + # Efficient slicing for contiguous frames + idx[t_axis] = frame_indices + else: + # Fancy indexing for non-contiguous + idx[t_axis] = list(indices) + data = dataset[tuple(idx)] + + if do != (1, 2, 0): + data = np.transpose(data, (do[2], do[0], do[1])) else: - # Fancy indexing for non-contiguous - data = dataset[indices, :, :] + # 4D dataset stored as (T, H, W, C); single-dataset case + if isinstance(frame_indices, slice): + data = dataset[frame_indices] + else: + data = dataset[list(indices)] + return np.asarray(data) output[:, :, :, ch_idx] = data diff --git a/src/pyflowreg/util/xcorr_prealignment.py b/src/pyflowreg/util/xcorr_prealignment.py index d26e288..d0accfa 100644 --- a/src/pyflowreg/util/xcorr_prealignment.py +++ b/src/pyflowreg/util/xcorr_prealignment.py @@ -31,7 +31,11 @@ def estimate_rigid_xcorr_2d( disambiguate : bool Whether to disambiguate sign of shift weight : np.ndarray or None - Channel weights for multi-channel images + Weights used to collapse multi-channel images to a single plane + before correlation. Either per-channel weights of shape (C,) or + spatial weights of shape (H, W, C) (a per-pixel weighted channel + mean is used; spatially constant weights are equivalent to the + per-channel form) Returns ------- @@ -41,10 +45,20 @@ def estimate_rigid_xcorr_2d( # Handle multi-channel images if ref_img.ndim == 3 and ref_img.shape[2] > 1: if weight is not None: - w = weight.reshape(-1).astype(np.float32) - w /= w.sum() - ref_img = np.tensordot(ref_img, w, axes=([2], [0])) - mov_img = np.tensordot(mov_img, w, axes=([2], [0])) + w = np.asarray(weight, dtype=np.float32) + if w.ndim >= 2: + # Spatial (H, W, C) weights (e.g. from reference + # preregistration): per-pixel weighted channel mean. + # Identical to the per-channel path when the weights are + # spatially constant. + wsum = w.sum(axis=-1) + 1e-12 + ref_img = (ref_img * w).sum(axis=-1) / wsum + mov_img = (mov_img * w).sum(axis=-1) / wsum + else: + w = w.reshape(-1) + w /= w.sum() + ref_img = np.tensordot(ref_img, w, axes=([2], [0])) + mov_img = np.tensordot(mov_img, w, axes=([2], [0])) else: ref_img = ref_img.mean(axis=2) mov_img = mov_img.mean(axis=2) diff --git a/tests/conftest.py b/tests/conftest.py index 6c64d5c..c437c4c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -109,19 +109,21 @@ def fast_of_options(temp_dir): @pytest.fixture(scope="function") def sequential_config(): """Create configuration for sequential executor.""" - return RegistrationConfig(n_jobs=1, verbose=True, parallelization="sequential") + return RegistrationConfig(n_jobs=1, verbose=False, parallelization="sequential") @pytest.fixture(scope="function") def threading_config(): """Create configuration for threading executor.""" - return RegistrationConfig(n_jobs=2, verbose=True, parallelization="threading") + return RegistrationConfig(n_jobs=2, verbose=False, parallelization="threading") @pytest.fixture(scope="function") def multiprocessing_config(): """Create configuration for multiprocessing executor.""" - return RegistrationConfig(n_jobs=2, verbose=True, parallelization="multiprocessing") + return RegistrationConfig( + n_jobs=2, verbose=False, parallelization="multiprocessing" + ) @pytest.fixture(scope="function") @@ -130,7 +132,7 @@ def auto_config(): return RegistrationConfig( n_jobs=2, batch_size=10, - verbose=True, + verbose=False, parallelization=None, # Auto-select ) @@ -139,7 +141,7 @@ def auto_config(): def executor_config(request): """Parametrized fixture to test all executor types.""" return RegistrationConfig( - n_jobs=2, batch_size=5, verbose=True, parallelization=request.param + n_jobs=2, batch_size=5, verbose=False, parallelization=request.param ) diff --git a/tests/core/test_diso_optical_flow.py b/tests/core/test_diso_optical_flow.py index 4b650c4..9014642 100644 --- a/tests/core/test_diso_optical_flow.py +++ b/tests/core/test_diso_optical_flow.py @@ -308,7 +308,7 @@ def test_diso_factory_with_params(self): assert flow.dtype == np.float32 def test_diso_factory_with_initial_flow(self): - """Test factory with initial flow.""" + """Test factory with initial flow passed as the pipeline keyword uv.""" diso_fn = _diso_factory() H, W = 64, 64 @@ -322,6 +322,83 @@ def test_diso_factory_with_initial_flow(self): assert flow.shape == (H, W, 2) assert flow.dtype == np.float32 + # The warm start must not be mutated in place (cv2 writes the + # result into the init buffer unless the wrapper copies it). + assert np.array_equal(w_init, np.ones((H, W, 2), dtype=np.float32)) + + +class _FakeDIS: + """Stands in for cv2.DISOpticalFlow to capture the init flow argument.""" + + def __init__(self): + self.received_inits = [] + + def calc(self, A, B, init): + self.received_inits.append(init) + if init is not None: + return init + return np.zeros(A.shape[:2] + (2,), dtype=np.float32) + + +class TestDisoOFWarmStart: + """DisoOF must honor the pipeline flow initialization keyword uv.""" + + def _diso_with_fake(self): + diso = DisoOF() + fake = _FakeDIS() + diso._dis = fake # _ensure() keeps an existing object + return diso, fake + + def test_call_uv_reaches_cv2_as_init(self): + """uv= is forwarded to cv2 DIS as the initial flow, as a copy.""" + diso, fake = self._diso_with_fake() + H, W = 32, 48 + fixed = np.random.rand(H, W).astype(np.float32) + moving = np.random.rand(H, W).astype(np.float32) + uv = np.full((H, W, 2), 1.5, dtype=np.float32) + + diso(fixed, moving, uv=uv) + + init = fake.received_inits[0] + assert init is not None + assert np.array_equal(init, uv) + assert init is not uv, "init must be a copy, not the caller's buffer" + + def test_call_w_takes_precedence_over_uv(self): + """Explicit w= wins when both w and uv are supplied.""" + diso, fake = self._diso_with_fake() + H, W = 32, 48 + fixed = np.random.rand(H, W).astype(np.float32) + moving = np.random.rand(H, W).astype(np.float32) + w = np.full((H, W, 2), 2.0, dtype=np.float32) + uv = np.full((H, W, 2), -3.0, dtype=np.float32) + + diso(fixed, moving, w=w, uv=uv) + + assert np.array_equal(fake.received_inits[0], w) + + def test_call_mismatched_uv_is_skipped(self): + """A warm start with the wrong spatial size is silently dropped.""" + diso, fake = self._diso_with_fake() + H, W = 32, 48 + fixed = np.random.rand(H, W).astype(np.float32) + moving = np.random.rand(H, W).astype(np.float32) + uv = np.ones((H + 2, W, 2), dtype=np.float32) + + diso(fixed, moving, uv=uv) + + assert fake.received_inits[0] is None + + def test_call_without_init_passes_none(self): + """No w/uv (e.g. FlowRegLive's first frame) runs a cold start.""" + diso, fake = self._diso_with_fake() + H, W = 32, 48 + fixed = np.random.rand(H, W).astype(np.float32) + moving = np.random.rand(H, W).astype(np.float32) + + diso(fixed, moving, uv=None) + + assert fake.received_inits[0] is None def test_diso_factory_with_weights(self): """Test factory with channel weights.""" diff --git a/tests/docs/test_docstring_mirrors.py b/tests/docs/test_docstring_mirrors.py index d904ec8..60c54e8 100644 --- a/tests/docs/test_docstring_mirrors.py +++ b/tests/docs/test_docstring_mirrors.py @@ -24,9 +24,9 @@ Speed adjustments versus the published examples (intent is unchanged): tiny synthetic inputs, ``quality_setting="fast"`` where the example used defaults, explicit small ``reference_frames`` (the default ``50:500`` -exceeds the tiny clips), and the sequential executor to avoid Windows -process-spawn overhead. Examples that explicitly teach -``quality_setting="balanced"`` keep it. +exceeds the tiny clips and would clamp to the last frame), and the +sequential executor to avoid Windows process-spawn overhead. Examples +that explicitly teach ``quality_setting="balanced"`` keep it. """ from pathlib import Path diff --git a/tests/docs/user_guide/test_multi_session.py b/tests/docs/user_guide/test_multi_session.py index 74569cd..874dc57 100644 --- a/tests/docs/user_guide/test_multi_session.py +++ b/tests/docs/user_guide/test_multi_session.py @@ -105,7 +105,12 @@ def test_session_pipeline_executes(self, materialize_session, snippet_runner): output_folders = ns["output_folders"] assert len(output_folders) == 3 for folder in output_folders: - assert (folder / "compensated.hdf5").exists() + # Stage 1 writes "compensated.HDF5"; the session module itself + # accepts both spellings (see stage1_compensate), and only + # case-insensitive filesystems collapse them. + assert (folder / "compensated.HDF5").exists() or ( + folder / "compensated.hdf5" + ).exists() assert (folder / "temporal_average.npy").exists() assert (folder / "idx.hdf").exists() diff --git a/tests/motion_correction/test_OF_options.py b/tests/motion_correction/test_OF_options.py index a377a71..e1ea086 100644 --- a/tests/motion_correction/test_OF_options.py +++ b/tests/motion_correction/test_OF_options.py @@ -112,6 +112,121 @@ def test_preregistration_creates_spatial_weights(self, tmp_path): assert reference.shape == (H, W, C) +class TestGetReferenceFrameIndexHandling: + """Clipping/deduplication of reference indices and cc forwarding.""" + + @staticmethod + def _write_video(tmp_path, T=6, H=16, W=24, C=1): + import tifffile + + frames = (np.random.rand(T, H, W, C) * 1000).astype(np.uint16) + video_path = tmp_path / "test_video.tif" + tifffile.imwrite(str(video_path), frames) + return video_path, frames + + @staticmethod + def _patch_compensate_arr(captured): + import importlib + from unittest.mock import patch + + # The pyflowreg.motion_correction package re-exports the + # compensate_arr *function*, shadowing the submodule attribute, so + # ``import ... as`` would bind the function. import_module returns + # the real submodule, which get_reference_frame imports from at + # call time. + compensate_arr_module = importlib.import_module( + "pyflowreg.motion_correction.compensate_arr" + ) + + def fake_compensate_arr(frames, reference, options=None, **kwargs): + captured["frames_shape"] = frames.shape + captured["options"] = options + T, H, W, _ = frames.shape + return frames, np.zeros((T, H, W, 2), dtype=np.float32) + + return patch.object( + compensate_arr_module, "compensate_arr", side_effect=fake_compensate_arr + ) + + def test_get_reference_frame_clips_and_dedupes_out_of_range_indices( + self, tmp_path, capsys + ): + """Clipped duplicate indices collapse to unique frames before prereg.""" + video_path, _ = self._write_video(tmp_path, T=6) + opts = OFOptions( + input_file=str(video_path), + reference_frames=list(range(2, 20)), # 2..5 valid, 6..19 clip to 5 + quality_setting=QualitySetting.FAST, + ) + reader = opts.get_video_reader() + + captured = {} + with self._patch_compensate_arr(captured): + opts.get_reference_frame(reader) + + # Unique frames 2, 3, 4, 5 — not 18 frames with 14 copies of frame 5. + assert captured["frames_shape"][0] == 4 + warning = capsys.readouterr().out + assert "exceed video length" in warning + assert "removing 14 duplicate indices" in warning + + def test_get_reference_frame_all_clipped_returns_single_frame(self, tmp_path): + """Default-style fully out-of-range lists skip preregistration.""" + video_path, frames = self._write_video(tmp_path, T=6) + opts = OFOptions( + input_file=str(video_path), + reference_frames=list(range(50, 500)), # all clip to frame 5 + quality_setting=QualitySetting.FAST, + ) + reader = opts.get_video_reader() + + captured = {} + with self._patch_compensate_arr(captured): + reference = opts.get_reference_frame(reader) + + assert "frames_shape" not in captured, "prereg must be skipped" + assert reference.shape == frames.shape[1:] + assert np.array_equal(reference, frames[5]) + + def test_get_reference_frame_preserves_in_range_duplicates(self, tmp_path): + """Intentional duplicates in fully in-range lists are kept.""" + video_path, _ = self._write_video(tmp_path, T=6) + opts = OFOptions( + input_file=str(video_path), + reference_frames=[1, 1, 2], + quality_setting=QualitySetting.FAST, + ) + reader = opts.get_video_reader() + + captured = {} + with self._patch_compensate_arr(captured): + opts.get_reference_frame(reader) + + assert captured["frames_shape"][0] == 3 + + def test_get_reference_frame_forwards_cc_settings(self, tmp_path): + """Preregistration options carry cc_initialization/cc_hw/cc_up.""" + video_path, _ = self._write_video(tmp_path, T=6) + opts = OFOptions( + input_file=str(video_path), + reference_frames=[0, 1, 2], + quality_setting=QualitySetting.FAST, + cc_initialization=True, + cc_hw=64, + cc_up=2, + ) + reader = opts.get_video_reader() + + captured = {} + with self._patch_compensate_arr(captured): + opts.get_reference_frame(reader) + + prereg_options = captured["options"] + assert prereg_options.cc_initialization is True + assert prereg_options.cc_hw == 64 + assert prereg_options.cc_up == 2 + + class TestAlphaValidation: """Test alpha parameter validation.""" diff --git a/tests/motion_correction/test_compensate_arr.py b/tests/motion_correction/test_compensate_arr.py index b92f452..2c77755 100644 --- a/tests/motion_correction/test_compensate_arr.py +++ b/tests/motion_correction/test_compensate_arr.py @@ -89,6 +89,18 @@ def test_multichannel_input(self): assert registered.shape == (T, H, W, C) assert flow.shape == (T, H, W, 2) + def test_single_frame_multichannel_input(self): + """A single (H,W,C) frame with a 3D reference mirrors its rank.""" + H, W, C = 24, 24, 2 + frame = np.random.rand(H, W, C).astype(np.float32) + reference = np.random.rand(H, W, C).astype(np.float32) + + registered, flow = compensate_arr(frame, reference) + + # Output rank mirrors the input rank: no leading time axis. + assert registered.shape == (H, W, C) + assert flow.shape == (H, W, 2) + class TestCompensateArrWithOptions: """Test compensate_arr with various OFOptions configurations.""" @@ -678,3 +690,20 @@ def test_options_not_modified(self): assert ( original_options.output_format == original_format ) # Should not be changed to ARRAY + + def test_backend_params_not_aliased_or_mutated(self): + """Neither the user's options nor the kwarg dict leak mutations.""" + T, H, W, C = 3, 16, 16, 1 + video = np.random.rand(T, H, W, C).astype(np.float32) + reference = np.mean(video, axis=0) + + original_options = OFOptions(quality_setting="fast") + assert original_options.backend_params == {} + + # The flowreg factory accepts and ignores extra keyword arguments. + kwarg_params = {"unused_param": 1} + compensate_arr(video, reference, original_options, backend_params=kwarg_params) + + # The caller's options and kwarg dict stay untouched. + assert original_options.backend_params == {} + assert kwarg_params == {"unused_param": 1} diff --git a/tests/motion_correction/test_compensate_recording.py b/tests/motion_correction/test_compensate_recording.py index 58afb26..e0cae68 100644 --- a/tests/motion_correction/test_compensate_recording.py +++ b/tests/motion_correction/test_compensate_recording.py @@ -35,6 +35,20 @@ def test_custom_config(self): assert config.verbose is True assert config.parallelization == "threading" + @pytest.mark.parametrize("verbose", [False, True]) + def test_verbose_gates_progress_output(self, fast_of_options, verbose, capsys): + """Progress messages print only when verbose=True (natural polarity).""" + config = RegistrationConfig( + n_jobs=1, verbose=verbose, parallelization="sequential" + ) + BatchMotionCorrector(fast_of_options, config) + + out = capsys.readouterr().out + if verbose: + assert "Using sequential executor" in out + else: + assert "Using sequential executor" not in out + class TestCompensateRecording: """Test the CompensateRecording class and executor system.""" @@ -121,7 +135,7 @@ def test_flow_params_include_constancy_assumption(self, tmp_path): iterations=2, ) config = RegistrationConfig( - n_jobs=1, verbose=True, parallelization="sequential" + n_jobs=1, verbose=False, parallelization="sequential" ) pipeline = BatchMotionCorrector(options, config) pipeline.weight = np.ones((8, 8, 1), dtype=np.float64) @@ -209,7 +223,7 @@ def test_compensate_recording_sequential(self, small_test_video, fast_of_options fast_of_options.buffer_size = 5 config = RegistrationConfig( - n_jobs=1, verbose=True, parallelization="sequential" + n_jobs=1, verbose=False, parallelization="sequential" ) # Test that pipeline can be created and configured correctly @@ -229,7 +243,7 @@ def test_compensate_recording_all_executors( fast_of_options.buffer_size = 3 config = RegistrationConfig( - n_jobs=2, verbose=True, parallelization=executor_name + n_jobs=2, verbose=False, parallelization=executor_name ) # Test executor selection by creating pipeline and checking executor type diff --git a/tests/motion_correction/test_flow_reg_live.py b/tests/motion_correction/test_flow_reg_live.py index 8f1ace0..26b6130 100644 --- a/tests/motion_correction/test_flow_reg_live.py +++ b/tests/motion_correction/test_flow_reg_live.py @@ -112,6 +112,68 @@ def test_set_reference_from_multiple_frames(self): assert flow_reg.reference_raw.shape == (H, W, C) assert flow_reg.reference_proc.shape == (H, W, C) + def test_set_reference_from_grayscale_stack(self): + """A 3D (T,H,W) array with last dim > 4 is a stack, not one frame.""" + flow_reg = FlowRegLive() + + T, H, W = 8, 24, 32 + frames = np.random.rand(T, H, W).astype(np.float32) + + flow_reg.set_reference(frames) + + # Preregistered and averaged to a single-channel (H,W,1) reference, + # not misread as one (T,H,W)="(H,W,C)" frame with W channels. + assert flow_reg.reference_raw.shape == (H, W, 1) + assert flow_reg.reference_proc.shape == (H, W, 1) + + def test_set_reference_norm_bounds_from_raw_reference(self): + """Normalization bounds come from the RAW reference's range.""" + flow_reg = FlowRegLive() + + H, W, C = 24, 24, 1 + # Raw-scale (uint16-like) reference with a known range + reference = (np.random.rand(H, W, C) * 40000 + 1000).astype(np.uint16) + flow_reg.set_reference(reference) + + assert flow_reg.norm_min == reference.min() + assert flow_reg.norm_max == reference.max() + + def test_call_handles_raw_scale_uint16_frames(self): + """Raw uint16 frames are normalized consistently with the reference.""" + flow_reg = FlowRegLive() + + H, W = 32, 32 + rng = np.random.default_rng(0) + base = (rng.random((H, W, 1)) * 40000 + 1000).astype(np.uint16) + flow_reg.set_reference(base) + + registered, flow = flow_reg(base.copy()) + + assert np.all(np.isfinite(registered)) + assert np.all(np.isfinite(flow)) + # Identical raw-scale frame vs reference: displacements stay small + # (with the old filtered-reference bounds the frame reached the + # solver at ~4e4 scale against a [0,1] reference). + assert np.abs(flow).max() < 2.0 + + def test_per_channel_sigmas_are_kept(self): + """Per-channel sigma rows are not collapsed to channel 0's.""" + options = OFOptions(sigma=[[1.0, 2.0, 0.5], [3.0, 4.0, 1.0]]) + flow_reg = FlowRegLive(options=options) + + assert flow_reg.sigma_2d.shape == (2, 2) + assert np.allclose(flow_reg.sigma_2d, [[1.0, 2.0], [3.0, 4.0]]) + assert flow_reg.sigma_t == 1.0 + + def test_call_takes_only_the_frame(self): + """__call__ no longer accepts the removed dead normalize kwarg.""" + flow_reg = FlowRegLive() + frame = np.random.rand(16, 16, 1).astype(np.float32) + flow_reg.set_reference(frame) + + with pytest.raises(TypeError): + flow_reg(frame, True) + def test_set_reference_from_buffer(self): """Test setting reference from internal buffer.""" flow_reg = FlowRegLive(reference_buffer_size=5) diff --git a/tests/util/io/test_tiff.py b/tests/util/io/test_tiff.py index c443918..7f2ced9 100644 --- a/tests/util/io/test_tiff.py +++ b/tests/util/io/test_tiff.py @@ -17,6 +17,28 @@ def _write_interleaved_pages_tiff(path, data): tw.write(data[t_idx, :, :, c_idx]) +def test_fresh_reader_has_batch_initializes(tmp_path): + """ + Regression test: has_batch() on a freshly constructed (lazy) reader + must initialize it, so `for batch in reader:` yields all frames + instead of terminating immediately with frame_count == 0. + """ + t_count, height, width = 6, 8, 6 + data = (np.random.rand(t_count, height, width) * 1000).astype(np.uint16) + tif_path = tmp_path / "stack.tif" + tifffile.imwrite(tif_path, data) + + reader = TIFFFileReader(str(tif_path), buffer_size=4) + try: + assert reader.has_batch() + batches = [batch for batch in reader] + frames = np.concatenate(batches, axis=0) + assert frames.shape == (t_count, height, width, 1) + np.testing.assert_array_equal(frames[..., 0], data) + finally: + reader.close() + + def test_read_batch_overreported_frame_count(tmp_path): """ Guard existing behavior: diff --git a/tests/util/test_xcorr_prealignment.py b/tests/util/test_xcorr_prealignment.py index 69c0319..7a1714e 100644 --- a/tests/util/test_xcorr_prealignment.py +++ b/tests/util/test_xcorr_prealignment.py @@ -136,6 +136,66 @@ def test_weighted_estimation(self): assert np.abs(estimated_shift[0] - true_dx_dy[0]) <= 1.0 assert np.abs(estimated_shift[1] - true_dx_dy[1]) <= 1.0 + @staticmethod + def _multichannel_pair(true_dx_dy): + """Structured 2-channel ref/mov pair shifted by true_dx_dy.""" + H, W, C = 64, 64, 2 + rng = np.random.default_rng(0) + ref = np.zeros((H, W, C), dtype=np.float32) + y, x = np.ogrid[:H, :W] + ref[..., 0] = np.exp(-((y - 20) ** 2 + (x - 30) ** 2) / 50) + ref[..., 1] = np.exp(-((y - 40) ** 2 + (x - 35) ** 2) / 40) + ref[15:25, 40:50, 0] = 1.0 + ref[35:45, 10:20, 1] = 0.8 + ref += rng.standard_normal((H, W, C)).astype(np.float32) * 0.05 + + mov = np.zeros_like(ref) + for c in range(C): + mov[..., c] = ndi_shift( + ref[..., c], + shift=(true_dx_dy[1], true_dx_dy[0]), + order=1, + mode="constant", + ) + return ref, mov + + def test_estimate_rigid_xcorr_2d_spatial_weight_runs(self): + """Spatial (H, W, C) weights (preregistration shape) are accepted.""" + true_dx_dy = np.array([3.0, -2.0], dtype=np.float32) + ref, mov = self._multichannel_pair(true_dx_dy) + H, W, C = ref.shape + + # Per-channel weights broadcast to the spatial (H, W, C) layout + # that get_reference_frame builds for preregistration options. + weight_spatial = np.zeros((H, W, C), dtype=np.float32) + weight_spatial[..., 0] = 0.7 + weight_spatial[..., 1] = 0.3 + + estimated_shift = estimate_rigid_xcorr_2d( + ref, mov, weight=weight_spatial, target_hw=(64, 64), up=1 + ) + + assert np.abs(estimated_shift[0] - true_dx_dy[0]) <= 1.0 + assert np.abs(estimated_shift[1] - true_dx_dy[1]) <= 1.0 + + def test_estimate_rigid_xcorr_2d_spatial_weight_matches_channel_weight(self): + """Spatially constant (H, W, C) weights equal the (C,) weight path.""" + true_dx_dy = np.array([3.0, -2.0], dtype=np.float32) + ref, mov = self._multichannel_pair(true_dx_dy) + H, W, C = ref.shape + + weight_channel = np.array([0.7, 0.3], dtype=np.float32) + weight_spatial = np.broadcast_to(weight_channel, (H, W, C)).copy() + + shift_channel = estimate_rigid_xcorr_2d( + ref, mov, weight=weight_channel, target_hw=(64, 64), up=10 + ) + shift_spatial = estimate_rigid_xcorr_2d( + ref, mov, weight=weight_spatial, target_hw=(64, 64), up=10 + ) + + np.testing.assert_allclose(shift_spatial, shift_channel, atol=1e-4) + class TestCCPrealignmentIntegration: """Test CC prealignment integration with motion compensation."""