diff --git a/.github/workflows/docs-test.yml b/.github/workflows/docs-test.yml index 9be727ad..e46248f2 100644 --- a/.github/workflows/docs-test.yml +++ b/.github/workflows/docs-test.yml @@ -22,10 +22,9 @@ jobs: cache: "pip" - name: Install dependencies - run: pip install -e .[docs] pytest-md + run: | + pip install pdm + pdm install --with docs - name: Build docs - working-directory: docs - env: - SPHINXOPTS: "-W --keep-going" - run: make html + run: pdm run docs-prod diff --git a/.gitignore b/.gitignore index 7ab1a9cb..ffae3ef1 100644 --- a/.gitignore +++ b/.gitignore @@ -165,4 +165,8 @@ user_data/* # The default output directory for the process commands. Not necessary but doesn't hurt to have. mio_process/* !user_data/.gitkeep +tests/data/stitch/*_timestamps.csv +tests/data/stitch/*_noise.csv +tests/data/stitch/*_scores.csv +tests/data/stitch/*_stitched* diff --git a/docs/api/models/dataset.md b/docs/api/models/dataset.md new file mode 100644 index 00000000..f507cb6b --- /dev/null +++ b/docs/api/models/dataset.md @@ -0,0 +1,7 @@ +# dataset + +```{eval-rst} +.. automodule:: mio.models.dataset + :members: + :undoc-members: +``` \ No newline at end of file diff --git a/docs/api/models/index.md b/docs/api/models/index.md index ae6c6cf5..107271ab 100644 --- a/docs/api/models/index.md +++ b/docs/api/models/index.md @@ -19,6 +19,7 @@ keep what is common common, and what is unique unique. buffer config data +dataset mixins models sdcard diff --git a/docs/conf.py b/docs/conf.py index e6796f4a..5c1f1305 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -81,6 +81,8 @@ # Mock imports for packages we don't have yet - this one is # for opal kelley stuff we need to figure out the licensing for autodoc_mock_imports = ["routine"] +autodoc_pydantic_model_show_json = False +autodoc_pydantic_model_show_json_error_strategy = "coerce" # todo todo_include_todos = True diff --git a/docs/meta/changelog.md b/docs/meta/changelog.md index 91d4bd5a..45553049 100644 --- a/docs/meta/changelog.md +++ b/docs/meta/changelog.md @@ -1,8 +1,8 @@ # Changelog -## Upcoming +## 0.10 -### *.* +### 0.10.0 #### CLI @@ -12,11 +12,32 @@ - `mio config open` to open the config in default text editor - [`#154`](https://github.com/miniscope/mio/pull/154) - add cli command for removing frames from video: - `mio process remove_frames` to remove frames by explicitly specified index from videos and metadata +- [`#155`](https://github.com/miniscope/mio/pull/155) - `mio process concat` - concatenate videos and metadata #### CI/CD - [`#157`](https://github.com/miniscope/mio/pull/157) - Add continuous deployment to PyPI +#### New features + +- [`#133`](https://github.com/miniscope/mio/pull/133) - {class}`~mio.models.dataset.Dataset` + organization - group recordings with their metadata, and group multiple recordings collected at the same time. +- [`#133`](https://github.com/miniscope/mio/pull/133), [`#155`](https://github.com/miniscope/mio/pull/155) + Noise-aware stitching: Given two recordings of the same data stream, + create a stitched version that picks the best frames from each of them +- [`#133`](https://github.com/miniscope/mio/pull/133), [`#155`](https://github.com/miniscope/mio/pull/155) + Alignment Maps - within a dataset, create an alignment map to align frames between recordings, + either by `frame_num` or by timestamps. +- preserve noise scoring metadata in `_scores.csv` and use to pick frames during stitching + +#### Perf + +- [`#155`](https://github.com/miniscope/mio/pull/155) - Vectorized black area detection + +#### Removed + +- [`#155`](https://github.com/miniscope/mio/pull/155) - Inter-frame mean squared error noise detection, unused. + ## 0.9 ### 0.9.0 - 2026-01-27 - Batch device update, NTP sync, driver import fix diff --git a/mio/cli/process.py b/mio/cli/process.py index 68994a7c..5d201c37 100644 --- a/mio/cli/process.py +++ b/mio/cli/process.py @@ -9,10 +9,12 @@ from mio.logging import init_logger from mio.models.dataset import Recording from mio.models.process import DenoiseConfig +from mio.process.stitch import concat_recordings from mio.process.stitch import stitch as run_stitch from mio.process.video import denoise as run_denoise from mio.process.video import remove_frames as run_remove_frames from mio.process.video import trim as run_trim +from mio.types import ConfigSource logger = init_logger("mio.cli.process") @@ -198,6 +200,42 @@ def remove_frames(input: str, output: str | None, frames: str, force: bool = Fal click.echo(f"Video written to {output} with frames {frames} removed from {input}") +@process.command() +@click.option( + "-i", + "--inputs", + required=True, + multiple=True, + type=click.Path(exists=True, dir_okay=False), + help="Paths to video files. Each requires a .csv with the same stem name.", +) +@click.option( + "-o", + "--output", + type=click.Path(dir_okay=False), + required=True, + help="Path to the output concatenated video file or directory. " + "If not specified, saves next to video with '_combined' suffix.", +) +def concat( + inputs: list[Path], + output: Path, +) -> None: + """ + Concatenate sequential recording segments from one DAQ into a single video. + + Use this to combine multiple segment files (e.g. long-2.avi, long-3.avi, ...) + from the same DAQ before stitching across DAQs. + """ + if len(inputs) < 2: + raise click.ClickException("Need at least 2 .avi files to concat") + recordings = [Recording.from_video(Path(p)) for p in inputs] + output = Path(output) + + click.echo(f"Concatenating {len(recordings)} segments...") + concat_recordings(recordings=recordings, output_video_path=output, progress=True) + + @process.command() @click.option( "-i", @@ -214,6 +252,13 @@ def remove_frames(input: str, output: str | None, frames: str, force: bool = Fal default=None, help="Directory for output videos and metadata. If none provided, same as the inputs.", ) +@click.option( + "-c", + "--config", + default=None, + help="A config id or path for a DenoiseConfig used to score frames if no noise score exists." + "If not provided, default config is used.", +) @click.option( "--debug-video", default=False, @@ -228,7 +273,11 @@ def remove_frames(input: str, output: str | None, frames: str, force: bool = Fal help="Overwrite any existing files", ) def stitch( - inputs: tuple, output: Path | None = None, debug_video: bool = False, force: bool = False + inputs: tuple, + output: Path | None = None, + config: ConfigSource | None = None, + debug_video: bool = False, + force: bool = False, ) -> None: """ Stitch multiple video recordings into one by selecting the best frame @@ -239,9 +288,20 @@ def stitch( if len(inputs) < 2: raise click.ClickException("At least 2 input videos are required for stitching.") + if config is not None: + denoise_config: DenoiseConfig = DenoiseConfig.from_any(config) + patch_config = denoise_config.noise_patch + else: + patch_config = None + recordings = [Recording.from_video(Path(p)) for p in inputs] stitched = run_stitch( - recordings, debug_video=debug_video, output_dir=output, progress=True, force=force + recordings, + debug_video=debug_video, + noise_config=patch_config, + output_dir=output, + progress=True, + force=force, ) click.echo(f"Stitched videos to {stitched.video.path}") @@ -309,13 +369,21 @@ def workflow( output_dir = Path(output).expanduser() output_dir.mkdir(parents=True, exist_ok=True) + denoise_config_parsed = DenoiseConfig.from_any(denoise_config) + if len(inputs) == 1: click.echo("Only one input video provided, skipping stitching") stitched_video = inputs[0] else: click.echo("Stitching videos...") recordings = [Recording.from_video(p) for p in inputs] - stitched = run_stitch(recordings, output_dir=output_dir, progress=True, force=force) + stitched = run_stitch( + recordings, + output_dir=output_dir, + noise_config=denoise_config_parsed, + progress=True, + force=force, + ) stitched_video = stitched.video.path if trim_start == 0 and trim_end == 0: @@ -331,7 +399,6 @@ def workflow( if trimmed.metadata is None: raise FileNotFoundError(f"No metadata csv found for video {trimmed_video}") - denoise_config_parsed = DenoiseConfig.from_any(denoise_config) final_video = run_denoise( trimmed_video, denoise_config_parsed, diff --git a/mio/data/config/process/denoise_example_mean_error.yml b/mio/data/config/process/denoise_calcium_imaging.yml similarity index 51% rename from mio/data/config/process/denoise_example_mean_error.yml rename to mio/data/config/process/denoise_calcium_imaging.yml index 728df692..4c8dec35 100644 --- a/mio/data/config/process/denoise_example_mean_error.yml +++ b/mio/data/config/process/denoise_calcium_imaging.yml @@ -1,45 +1,36 @@ -id: denoise_example_mean_error -mio_model: mio.models.process.DenoiseConfig -mio_version: 0.6.1 -noise_patch: - enable: true - method: [mean_error] - mean_error_config: - threshold: 40 - device_config_id: wireless-200px - buffer_split: 8 - comparison_unit: 1000 - diff_multiply: 1 - gradient_config: - threshold: 20 - black_area_config: - consecutive_threshold: 5 - value_threshold: 16 - output_result: true - output_noise_patch: true - output_diff: true - output_noisy_frames: true -frequency_masking: - id: frequency_masking_example_mean_error - mio_model: mio.models.process.FrequencyMaskingConfig - mio_version: 0.6.1 - enable: true - spatial_LPF_cutoff_radius: 15 - vertical_BEF_cutoff: 2 - horizontal_BEF_cutoff: 0 - output_mask: true - output_result: true - output_freq_domain: true -minimum_projection: - enable: true - normalize: true - output_result: true - output_min_proj: true -interactive_display: - show_videos: true - start_frame: 40 - end_frame: 140 - display_freq_mask: true -end_frame: -1 #-1 means all frames -output_result: true -output_dir: user_data/output \ No newline at end of file +id: denoise_calcium_imaging +mio_model: mio.models.process.DenoiseConfig +mio_version: 0.6.1 +noise_patch: + enable: true + method: [gradient, black_area] + gradient_config: + threshold: 20 + black_area_config: + consecutive_threshold: 30 + value_threshold: 0 + min_rows: 10 + output_result: true + output_noise_patch: true + output_noisy_frames: true +frequency_masking: + id: frequency_masking_calcium_imaging + mio_model: mio.models.process.FrequencyMaskingConfig + mio_version: 0.6.1 + enable: true + cast_float32: true + spatial_LPF_cutoff_radius: 15 + vertical_BEF_cutoff: 2 + horizontal_BEF_cutoff: 0 + output_result: true +minimum_projection: + enable: true + normalize: true + output_result: true +interactive_display: + show_videos: false + start_frame: 0 + end_frame: 100 + display_freq_mask: false +end_frame: -1 +output_result: true diff --git a/mio/data/config/process/denoise_example.yml b/mio/data/config/process/denoise_example.yml index 24bf075e..0c7aabef 100644 --- a/mio/data/config/process/denoise_example.yml +++ b/mio/data/config/process/denoise_example.yml @@ -4,12 +4,6 @@ mio_version: 0.6.1 noise_patch: enable: true method: [gradient, black_area] - mean_error_config: - threshold: 40 - device_config_id: wireless-200px - buffer_split: 8 - comparison_unit: 1000 - diff_multiply: 1 gradient_config: threshold: 20 black_area_config: @@ -17,7 +11,6 @@ noise_patch: value_threshold: 0 output_result: true output_noise_patch: true - output_diff: true output_noisy_frames: true frequency_masking: id: frequency_masking_example diff --git a/mio/data/config/process/denoise_patchonly.yml b/mio/data/config/process/denoise_patchonly.yml index 44db9145..d88e34d4 100644 --- a/mio/data/config/process/denoise_patchonly.yml +++ b/mio/data/config/process/denoise_patchonly.yml @@ -6,12 +6,6 @@ noise_patch: method: - gradient - black_area - mean_error_config: - threshold: 40 - device_config_id: wireless-200px - buffer_split: 8 - comparison_unit: 1000 - diff_multiply: 1 gradient_config: threshold: 20 black_area_config: @@ -19,7 +13,6 @@ noise_patch: value_threshold: 0 output_result: true output_noise_patch: true - output_diff: false output_noisy_frames: true frequency_masking: id: frequency_masking_example diff --git a/mio/models/dataset.py b/mio/models/dataset.py index 24ba6c2a..b72b71dc 100644 --- a/mio/models/dataset.py +++ b/mio/models/dataset.py @@ -74,6 +74,7 @@ ) from mio.models import MiniscopeIOModel +from mio.models.process import NoisePatchConfig from mio.utils import _format_ranges if sys.version_info < (3, 11): @@ -107,6 +108,8 @@ class RecordingPaths(TypedDict): """{stem}.csv""" timestamps: Path """{stem}_timestamps.csv""" + noise: Path + """{stem}_noise.csv""" binary: Path """{stem}.bin""" @@ -117,6 +120,7 @@ def paths_from_video(video: Path) -> RecordingPaths: video=video, metadata=video.with_suffix(".csv"), timestamps=video.with_name(video.stem + "_timestamps.csv"), + noise=video.with_name(video.stem + "_noise.csv"), binary=video.with_suffix(".bin"), ) @@ -138,6 +142,10 @@ class Recording(MiniscopeIOModel): When instantiating a recording, if a metadata file exists but timestamps do not, they are automatically generated. """ + noise: pd.DataFrame | None = None + """ + Framewise noise measurements (created with :meth:`score_noise` ). + """ binary: Path | None = None """Path to any raw binary version of the data in the video""" derived_from: RecordingDerivation | None = None @@ -159,6 +167,29 @@ def from_video(cls, path: Path) -> "RecordingUnion": else: return RawVideoRecording(name=path.stem, video=path) + def score_noise( + self, config: NoisePatchConfig | None = None, progress: bool = False, force: bool = False + ) -> pd.DataFrame: + """ + Score the noise level in each frame with :func:`.score_noise`, + saving as a csv with `{name}_noise.csv` + """ + + from mio.process.video import score_noise + + if config is None: + config = NoisePatchConfig() + if not force: + if self.noise is not None: + return self.noise + elif self.paths["noise"].exists(): + self.noise = pd.read_csv(self.paths["noise"]) + return self.noise + + self.noise = score_noise(self, config, progress=progress) + self.noise.to_csv(self.paths["noise"], index=False) + return self.noise + @model_validator(mode="before") @classmethod def _load_csvs(cls, v: dict) -> dict: diff --git a/mio/models/process.py b/mio/models/process.py index 3940122c..63fb02ea 100644 --- a/mio/models/process.py +++ b/mio/models/process.py @@ -8,7 +8,6 @@ from mio.models import MiniscopeConfig from mio.models.mixins import ConfigYAMLMixin -from mio.models.stream import StreamDevConfig class MinimumProjectionConfig(BaseModel): @@ -38,52 +37,13 @@ class MinimumProjectionConfig(BaseModel): ) -class MSEDetectorConfig(BaseModel): - """ - Configraiton for detecting invalid frames based on mean squared error. - """ - - threshold: float = Field( - ..., - description="Threshold for detecting invalid frames based on mean squared error.", - ) - device_config_id: str | None = Field( - default=None, - description="ID of the stream device configuration used for aquiring the video." - "This is used in the mean_error method to compare frames" - " in the units of data transfer buffers.", - ) - buffer_split: int = Field( - default=1, - description="Number of splits to make in the buffer when detecting noisy areas." - "This further splits the buffer into smaller patches to detect small noisy areas." - "This is used in the mean_error method.", - ) - diff_multiply: int = Field( - default=1, - description="Multiplier for visualizing the diff between the current and previous frame.", - ) - - _device_config: StreamDevConfig | None = None - - @property - def device_config(self) -> StreamDevConfig: - """ - Get the device configuration based on the device_config_id. - This is used in the mean_error method to compare frames in the units of data buffers. - """ - if self._device_config is None: - self._device_config = StreamDevConfig.from_any(self.device_config_id) - return self._device_config - - class GradientDetectorConfig(BaseModel): """ - Configraiton for detecting invalid frames based on gradient. + Configuration for detecting invalid frames based on gradient. """ threshold: float = Field( - ..., + default=20, description="Threshold for detecting invalid frames based on gradient.", ) @@ -101,6 +61,12 @@ class BlackAreaDetectorConfig(BaseModel): default=0, description="Pixel intensity value below which a pixel is considered 'black'.", ) + min_rows: int = Field( + default=1, + description="Minimum number of flagged rows required to mark the frame as invalid. " + "Default of 1 preserves original behavior. For calcium imaging, values around 10 " + "reduce false positives from naturally dark regions.", + ) class NoisePatchConfig(BaseModel): @@ -113,24 +79,18 @@ class NoisePatchConfig(BaseModel): default=True, description="Enable patch based noise handling.", ) - method: list[Literal["mean_error", "gradient", "black_area"]] = Field( - default="gradient", + method: list[Literal["gradient", "black_area"]] = Field( + default_factory=lambda: ["gradient", "black_area"], description="Method for detecting noise." "gradient: Detection based on the gradient of the frame row." - "mean_error: Detection based on the mean error with the same row of the previous frame." "black_area: Detection based on the number of consecutive black pixels in a row.", ) - mean_error_config: MSEDetectorConfig | None = Field( - default=None, - description="Configuration for detecting invalid frames based on mean squared error." - " Any positive value or zero is valid.", - ) - gradient_config: GradientDetectorConfig | None = Field( - default=None, + gradient_config: GradientDetectorConfig = Field( + default_factory=GradientDetectorConfig, description="Configuration for detecting invalid frames based on gradient.", ) - black_area_config: BlackAreaDetectorConfig | None = Field( - default=None, + black_area_config: BlackAreaDetectorConfig = Field( + default_factory=BlackAreaDetectorConfig, description="Configuration for detecting invalid frames based on black area.", ) output_result: bool = Field( @@ -142,12 +102,6 @@ class NoisePatchConfig(BaseModel): description="Output the noise patch video" "This highlights the noisy areas found in the video stream.", ) - output_diff: bool = Field( - default=False, - description="Output the diff video stream." - "The diff video stream shows the difference between the current and previous frame." - "This is used in the mean_error method.", - ) output_noisy_frames: bool = Field( default=True, description="Output the stack of noisy frames as an independent video stream.", diff --git a/mio/process/frame_helper.py b/mio/process/frame_helper.py index ca76485b..1dc095a5 100644 --- a/mio/process/frame_helper.py +++ b/mio/process/frame_helper.py @@ -2,6 +2,9 @@ This module contains a helper class for frame operations. """ +from __future__ import annotations + +import sys from abc import abstractmethod import cv2 @@ -12,13 +15,37 @@ BlackAreaDetectorConfig, FrequencyMaskingConfig, GradientDetectorConfig, - MSEDetectorConfig, NoisePatchConfig, ) +if sys.version_info < (3, 11): + from typing_extensions import TypedDict +else: + from typing import TypedDict + logger = init_logger("frame_helper") +class Detectors(TypedDict, total=False): + """Map between shortnames and detector class instances""" + + black_area: BlackAreaDetector + gradient: GradientNoiseDetector + + +def make_detectors(config: NoisePatchConfig) -> Detectors: + """Make detector classes from a config""" + detectors = {} + for method in config.method: + if method == "gradient": + detectors[method] = GradientNoiseDetector(config.gradient_config) + elif method == "black_area": + detectors[method] = BlackAreaDetector(config.black_area_config) + else: + raise ValueError(f"Unknown method {method}") + return detectors + + class BaseSingleFrameHelper: """ Base class for single frame operations. @@ -78,19 +105,7 @@ def __init__(self, noise_patch_config: NoisePatchConfig): if noise_patch_config.method is None: raise ValueError("No noise detection methods provided") self.methods = noise_patch_config.method - - if "mean_error" in self.methods: - if noise_patch_config.mean_error_config is None: - raise ValueError("Mean error config must be provided for mean error detection") - self.mse_detector = MSENoiseDetector(noise_patch_config.mean_error_config) - if "gradient" in self.methods: - if noise_patch_config.gradient_config is None: - raise ValueError("Gradient config must be provided for gradient detection") - self.gradient_detector = GradientNoiseDetector(noise_patch_config.gradient_config) - if "black_area" in self.methods: - if noise_patch_config.black_area_config is None: - raise ValueError("Black area config must be provided for black area detection") - self.black_detector = BlackAreaDetector(noise_patch_config.black_area_config) + self.detectors = make_detectors(noise_patch_config) def find_invalid_area(self, frame: np.ndarray) -> tuple[bool, np.ndarray]: """ @@ -104,19 +119,8 @@ def find_invalid_area(self, frame: np.ndarray) -> tuple[bool, np.ndarray]: """ noisy_flag = False combined_noisy_area = np.zeros_like(frame, dtype=np.uint8) - - if "mean_error" in self.methods: - noisy, noisy_area = self.mse_detector.find_invalid_area(frame) - combined_noisy_area = np.maximum(combined_noisy_area, noisy_area) - noisy_flag = noisy_flag or noisy - - if "gradient" in self.methods: - noisy, noisy_area = self.gradient_detector.find_invalid_area(frame) - combined_noisy_area = np.maximum(combined_noisy_area, noisy_area) - noisy_flag = noisy_flag or noisy - - if "black_area" in self.methods: - noisy, noisy_area = self.black_detector.find_invalid_area(frame) + for detector in self.detectors.values(): + noisy, noisy_area = detector.find_invalid_area(frame) combined_noisy_area = np.maximum(combined_noisy_area, noisy_area) noisy_flag = noisy_flag or noisy @@ -214,130 +218,36 @@ def _detect_black_pixels( current_frame: np.ndarray, ) -> tuple[bool, np.ndarray]: """ - Detect black-out noise by checking for black pixels (value 0) over rows of pixels. + Detect black-out noise by checking for consecutive black pixels per row. - Returns: - Tuple[bool, np.ndarray]: A boolean indicating if the frame is corrupted and noise mask. - """ - height, width = current_frame.shape - noisy_mask = np.zeros_like(current_frame, dtype=np.uint8) - - # Read values from YAML config - consecutive_threshold = ( - self.config.consecutive_threshold - ) # How many consecutive pixels must be black - black_pixel_value_threshold = ( - self.config.value_threshold - ) # Max pixel value considered "black" - - logger.debug(f"Using black pixel threshold: <= {black_pixel_value_threshold}") - logger.debug(f"Consecutive black pixel threshold: {consecutive_threshold}") - - frame_is_noisy = False # Track if frame should be discarded - - for y in range(height): - row = current_frame[y, :] # Extract row - consecutive_count = 0 # Counter for consecutive black pixels - - for x in range(width): - if row[x] <= black_pixel_value_threshold: # Check if pixel is "black" - consecutive_count += 1 - else: - consecutive_count = 0 # Reset if a non-black pixel is found - - # If we exceed the allowed threshold of consecutive black pixels, discard the frame - if consecutive_count >= consecutive_threshold: - logger.debug( - f"Frame noisy due to {consecutive_count} consecutive black pixels " - f"in row {y}." - ) - noisy_mask[y, :] = 1 # Mark row as noisy - frame_is_noisy = True - break # No need to check further in this row - - return frame_is_noisy, noisy_mask - - -class MSENoiseDetector(BaseSingleFrameHelper): - """ - Helper class for mean squared error noise detection. - """ - - def __init__(self, config: MSEDetectorConfig): - """ - Initialize the MeanErrorNoiseDetectionHelper object. - - Parameters: - threshold (float): The threshold for noise detection. + Uses vectorized numpy (cumulative sum sliding window) instead of + pixel-by-pixel Python loops for ~100x speedup. Returns: - MeanErrorNoiseDetectionHelper: A MeanErrorNoiseDetectionHelper object. - """ - self.config = config - self.previous_frame = None - - def register_previous_frame(self, previous_frame: np.ndarray) -> None: - """ - Register the previous frame for mean error calculation. - - Parameters: - previous_frame (np.ndarray): The previous frame to compare against. - """ - self.previous_frame = previous_frame - - def find_invalid_area(self, frame: np.ndarray) -> tuple[bool, np.ndarray]: - """ - Process a single frame and verify if it is valid. - - Parameters: - frame (np.ndarray): The frame to process. - - Returns: - Tuple[bool, np.ndarray]: A boolean indicating if the frame is valid - and the processed frame. - """ - if self.previous_frame is None: - self.previous_frame = frame - return False, np.zeros_like(frame, dtype=np.uint8) - noisy, mask = self._detect_with_mean_error(frame) - return noisy, mask - - def _detect_with_mean_error(self, current_frame: np.ndarray) -> tuple[bool, np.ndarray]: - """ - Detect noise using mean error between current and previous frames. - - Returns: - Tuple[bool, np.ndarray]: A boolean indicating if the frame is noisy and the noise mask. + Tuple[bool, np.ndarray]: A boolean indicating if the frame is corrupted and noise mask. """ - if self.previous_frame is None: - return False, np.zeros_like(current_frame, dtype=np.uint8) - - current_flat = current_frame.astype(np.int16).flatten() - previous_flat = self.previous_frame.astype(np.int16).flatten() + consecutive_threshold = self.config.consecutive_threshold + black_pixel_value_threshold = self.config.value_threshold - buffer_indices = FrameSplitter.get_buffer_shape( - current_frame.shape[1], current_frame.shape[0], self.config.device_config.px_per_buffer - ) + [ - current_frame.size - ] # Ensure final boundary is included + # Boolean mask of "black" pixels, then cumsum-based sliding window + black_mask = (current_frame <= black_pixel_value_threshold).astype(np.float32) + cs = np.cumsum(black_mask, axis=1) - noisy_mask = np.ones_like(current_flat, dtype=np.uint8) - has_noise = False + if current_frame.shape[1] >= consecutive_threshold: + # Sliding window: sum of `consecutive_threshold` consecutive pixels + run_sum = cs[:, consecutive_threshold:] - cs[:, :-consecutive_threshold] + # A row has a run if any window sums to exactly consecutive_threshold + # (meaning all pixels in that window were black) + row_has_run = np.any(run_sum >= consecutive_threshold, axis=1) + else: + row_has_run = np.zeros(current_frame.shape[0], dtype=bool) - for start_idx, end_idx in zip(buffer_indices[:-1], buffer_indices[1:]): - for sub_start in range( - end_idx - self.config.buffer_split, start_idx, -self.config.buffer_split - ): - mean_error = np.mean( - np.abs(current_flat[sub_start:end_idx] - previous_flat[sub_start:end_idx]) - ) - - if mean_error > self.config.threshold: - noisy_mask[sub_start:end_idx] = 0 - has_noise = True - break + noisy_mask = np.zeros_like(current_frame, dtype=np.uint8) + noisy_mask[row_has_run, :] = 1 - return has_noise, noisy_mask.reshape(current_frame.shape) + noisy_row_count = int(np.sum(row_has_run)) + frame_is_noisy = noisy_row_count >= self.config.min_rows + return frame_is_noisy, noisy_mask class FrequencyMaskHelper(BaseSingleFrameHelper): diff --git a/mio/process/stitch.py b/mio/process/stitch.py index 43f09558..60a590e7 100644 --- a/mio/process/stitch.py +++ b/mio/process/stitch.py @@ -1,25 +1,28 @@ """ -Buffer-wise stitching of multiple data streams based on device timestamps. +Buffer-wise stitching and concatenation of multiple data streams. This module combines multiple recordings (AVI video + metadata CSV) by selecting the best buffers from each stream using gradient noise detection. +It also provides concatenation of sequential recording segments from the same DAQ. This is still hardcoded around the StreamDevConfig metadata fields. """ from __future__ import annotations from dataclasses import dataclass, field +from functools import partial from pathlib import Path import cv2 import numpy as np import pandas as pd from pydantic import BaseModel -from tqdm import tqdm +from tqdm import tqdm, trange from mio.io import BufferedCSVWriter, VideoWriter from mio.logging import init_logger -from mio.models.dataset import Dataset, Recording, StitchedRecording +from mio.models.dataset import Dataset, Recording, StitchedRecording, paths_from_video +from mio.models.process import NoisePatchConfig logger = init_logger(name="stitch") @@ -28,24 +31,43 @@ def align(recordings: list[Recording]) -> pd.DataFrame: """ Create an alignment map by frame index. - Note that this **does not** align by timestamp! - it assumes that there is some ``frame_num`` in the metadata col for each of the recordings - that comes from some common device. + Note that this **is not** a general alignment method yet - + this is specialized to the case of stitching two recordings of the same underlying data source, + as is done when we record multiple FPGA sensors in the miniscope zero. + Please raise an issue if you need a general frame alignment mechanisms. + + We have two kinds of alignment, depending on the structure of the metadata: + + * If all the recordings have continuously incrementing `frame_num`s, + we align by the ``frame_num``. + The `frame_num` is given by the device, and is the same across recordings, + even if they start at different times (and capture different ranges of frame nums). + This is an **outer join**, keeping all frames + * If the recordings have *discontinuous* ``frame_num`` s, + e.g. if the device was restarted during acquisition, we align by the acquisition timestamp. + This assumes that the system times are closely matching + (specifically, more closely than the interval between successive frames in the recording). + This is an **inner join**, where we only keep frames where we can align timestamps. """ - metadatas: dict[str, pd.DataFrame] = {r.name: r.metadata for r in recordings} - if not all(isinstance(m, pd.DataFrame) for m in metadatas.values()): + if not all(isinstance(r.metadata, pd.DataFrame) for r in recordings): raise ValueError("All recordings must have metadata csvs to align them") if not all( - "frame_num" in m.columns and "reconstructed_frame_index" in m.columns - for m in metadatas.values() + "frame_num" in r.metadata.columns and "reconstructed_frame_index" in r.metadata.columns + for r in recordings ): raise ValueError("All recordings must have frame_num and reconstructed_frame_index columns") - # find the full set of frames in all the recordings - frame_set = set() - for m in metadatas.values(): - frame_set |= set(m["frame_num"]) + if not any(_has_discontinuous_runs(r.metadata["frame_num"]) for r in recordings): + logger.debug("Using frame-num based alignment") + return _align_by_frame(recordings) + else: + logger.debug("Using time-based alignment") + return _align_by_time(recordings) + +def _align_by_frame(recordings: list[Recording]) -> pd.DataFrame: + """Align metadata by the frame_num column""" + metadatas: dict[str, pd.DataFrame] = {r.name: r.metadata for r in recordings} # aggregate mappings from frame nums to the reconstructed frame index frame_maps = { name: df[["frame_num", "reconstructed_frame_index"]] @@ -71,8 +93,62 @@ def align(recordings: list[Recording]) -> pd.DataFrame: return aligned +def _align_by_time(recordings: list[Recording]) -> pd.DataFrame: + """ + Align by the nearest unix timestamp. + + Use the mean of the timestamps from the buffers to get frames that have the most overlap. + + This could be made an outer join by just keeping the leading and trailing rows, + and filtering rows with NaNs in the interior regions of buffer_recv_unix_time_x and y + but leaving as inner for now to match existing timestamp match fn. + + the inner join functions like "when both frames mutually pick each other as their closest frame" + which filters blippy frames that are very short. + + I **think** but have not tested that doing this triple merge method is faster + than nested iteration, esp for longer recordings, since these are all vector ops. + """ + metadatas: dict[str, pd.DataFrame] = {r.name: r.metadata for r in recordings} + time_maps = { + name: df.groupby("reconstructed_frame_index")["buffer_recv_unix_time"].mean().reset_index() + for name, df in metadatas.items() + } + + # inner join on closest mean timestamp value + names = sorted(time_maps.keys()) + last_name = names.pop(0) + aligned = time_maps[last_name].copy().rename(columns={"reconstructed_frame_index": last_name}) + for name in names: + # merge left and right, then take the inner match + left = pd.merge_asof( + aligned, time_maps[name], on="buffer_recv_unix_time", direction="nearest" + ) + right = pd.merge_asof( + time_maps[name], aligned, on="buffer_recv_unix_time", direction="nearest" + ) + left.rename(columns={"reconstructed_frame_index": name}, inplace=True) + right.rename(columns={"reconstructed_frame_index": name}, inplace=True) + + # merge on the frame indexes from the left and right - + # align when both sides agree they are the closest, + # dropping extras from glitches/sampling rate differences + aligned = pd.merge(left, right, "inner", on=[last_name, name]) + + # keep the left's times, keeping them anchored rather than wandering in each recording + aligned = aligned[[c for c in aligned.columns if c != "buffer_recv_unix_time_y"]] + aligned.rename(columns={"buffer_recv_unix_time_x": "buffer_recv_unix_time"}, inplace=True) + last_name = name + + aligned = aligned.astype({k: "Int64" for k in metadatas}) + # popping the index gives us the 'index' column + aligned = aligned.reset_index() + return aligned + + def stitch( recordings: list[Recording], + noise_config: NoisePatchConfig | None = None, dataset: Dataset | None = None, debug_video: bool = False, output_dir: Path | None = None, @@ -89,10 +165,25 @@ def stitch( It does not handle stitching or aligning videos that were recorded with *different* devices, for that use the :attr:`.Dataset.alignment_map` which aligns simultaneous frames in different recordings. + + Args: + recordings (list[Recording]): List of recordings to stitch. + noise_config (NoisePatchConfig | None): Configuration used for scoring + noise per frame (with :func:`.score_noise` ). If None, use defaults + dataset (Dataset | None): existing dataset, e.g. with existing alignment mapping + output_dir: (Path | None): where to write stitched video and metadata, + if None, same as recording directory + progress (bool): Show a progress bar. Default ``False`` + force (bool): Overwrite existing stitched video and metadata CSV files """ if len(recordings) != 2: raise NotImplementedError("Only stitching two videos simultaneously is supported!") + # ensure that the recordings have noise scores + # (does not recompute if they already exist) + for rec in recordings: + rec.score_noise(config=noise_config, progress=progress, force=force) + if dataset is None: dataset = Dataset.from_recordings(recordings) output_dir = dataset.path if output_dir is None else Path(output_dir) @@ -136,6 +227,10 @@ def stitch( frames.append(np.zeros(rec.video.shape[1:], dtype=np.uint8)) continue buffer_rows = rec.metadata[rec.metadata["reconstructed_frame_index"] == row[name]] + noise_row = rec.noise[rec.noise["reconstructed_frame_index"] == row[name]].iloc[0] + black_pixels = int(noise_row["black_area"]) if "black_area" in noise_row else 0 + noisy_pixels = int(noise_row["noisy_area"]) if "noisy_area" in noise_row else 0 + frames.append(rec.video[int(row[name])]) candidates.append( CandidateFrame( @@ -143,10 +238,12 @@ def stitch( frame=frames[-1], num_buffers=len(buffer_rows), sum_black_padding=int(buffer_rows["black_padding_px"].fillna(0).sum()), + black_pixels=black_pixels, + noisy_pixels=noisy_pixels, metadata_rows=buffer_rows, ) ) - result = _select_best_candidate(candidates, row["index"], row["frame_num"]) + result = _select_best_candidate(candidates, row["index"], row.get("frame_num")) selected = [c for c in candidates if c.recording.name == result.selected_video][0] if debug_video_writer is not None: debug_frame = np.zeros_like(frames[0], dtype=np.uint8) @@ -178,6 +275,8 @@ class CandidateFrame: frame: np.ndarray num_buffers: int sum_black_padding: int + black_pixels: int + noisy_pixels: int metadata_rows: pd.DataFrame _edge_score: float | None = field(default=None, repr=False) @@ -193,7 +292,9 @@ def metadata_score(self) -> tuple[int, int]: """Higher is better: more buffers, less black padding. A bit overkill but left this for future extension. """ - return (self.num_buffers, -self.sum_black_padding) + # To discuss - we are probably double counting padding and missing buffers, + # but keeping similar to existing method until we can decide what we want here -jls + return (self.num_buffers, -self.sum_black_padding - self.black_pixels - self.noisy_pixels) def _score_edges(frame: np.ndarray) -> float: @@ -211,13 +312,17 @@ class StitchRecord(BaseModel): """ index: int - frame_num: int + frame_num: int | None = None selected_video: str compare_video: str | None = None selected_num_buffers: int selected_black_padding: int + selected_black_pixels: int + selected_noisy_pixels: int compare_num_buffers: int | None = None compare_black_padding: int | None = None + compare_black_pixels: int | None = None + compare_noisy_pixels: int | None = None selected_edge_score: float | None = None compare_edge_score: float | None = None @@ -228,7 +333,7 @@ def header(cls) -> list[str]: def _select_best_candidate( - candidates: list[CandidateFrame], index: int, frame_num: int + candidates: list[CandidateFrame], index: int, frame_num: int | None = None ) -> StitchRecord: """ Pick the best candidate using metadata scoring with edge-score tiebreak. @@ -246,6 +351,8 @@ def _select_best_candidate( selected_video=candidates[0].recording.name, selected_num_buffers=candidates[0].num_buffers, selected_black_padding=candidates[0].sum_black_padding, + selected_black_pixels=candidates[0].black_pixels, + selected_noisy_pixels=candidates[0].noisy_pixels, **kwargs, ) @@ -269,8 +376,112 @@ def _select_best_candidate( selected_video=selected.recording.name, selected_num_buffers=selected.num_buffers, selected_black_padding=selected.sum_black_padding, + selected_black_pixels=selected.black_pixels, + selected_noisy_pixels=selected.noisy_pixels, compare_video=other.recording.name, compare_num_buffers=other.num_buffers, compare_black_padding=other.sum_black_padding, + compare_black_pixels=other.black_pixels, + compare_noisy_pixels=other.noisy_pixels, **kwargs, ) + + +def concat_recordings( + recordings: list[Recording], output_video_path: Path, progress: bool = False +) -> Recording: + """Concatenate sequential recording segments into a single video + CSV. + + Each recording's frames are appended in order. The CSV metadata is merged + with ``reconstructed_frame_index`` renumbered to be contiguous across all + segments. + + Parameters + ---------- + recordings : list[RecordingData] + Ordered list of recording segments to concatenate. + output_video_path : Path + Path for the combined output AVI. + progress : bool + Show a progress bar + """ + fps = int(recordings[0].video.video.get(cv2.CAP_PROP_FPS)) + video_writer = VideoWriter(path=output_video_path, fps=fps) + metadata_parts: list[pd.DataFrame] = [] + rfi_offset = 0 + total_frames = 0 + + recs = ( + tqdm(enumerate(recordings), desc="Concatenating recordings", position=0) + if progress + else enumerate(recordings) + ) + frame_iter_cls = partial(trange, position=1) if progress else range + try: + for i, rec in recs: + # Copy all video frames + seg_frames = 0 + total_frames = rec.video.shape[0] + + for n in frame_iter_cls(total_frames): + frame = rec.video[n] + video_writer.write_frame(frame) + seg_frames += 1 + + # Offset reconstructed_frame_index in metadata + df = rec.metadata.copy() + max_rfi = int(df["reconstructed_frame_index"].max()) + df["reconstructed_frame_index"] = df["reconstructed_frame_index"] + rfi_offset + metadata_parts.append(df) + + logger.debug( + "Segment %s: %s — %s frames, rfi_offset=%s", + i, + rec.video.path.name, + seg_frames, + rfi_offset, + ) + rfi_offset += max_rfi + 1 + total_frames += seg_frames + finally: + video_writer.close() + if progress: + recs.close() + + combined_df = pd.concat(metadata_parts, ignore_index=True) + combined_df.to_csv(paths_from_video(output_video_path)["metadata"], index=False) + + logger.debug( + "Concat completed: %s frames from %s segments -> %s", + total_frames, + len(recordings), + output_video_path, + ) + return Recording.from_video(output_video_path) + + +def _has_discontinuous_runs(series: pd.Series) -> bool: + """ + Check if a metadata series has multiple discontinuous series of values: + e.g. when acquiring frames and the counter is reset. + + Ignores single-row discontinuities like e.g. from a single buffer having an incorrect frame_num + """ + # we need the initial NaN for alignment below, so don't drop it yet - + # filtering NaNs is presumably cheaper than diffing + diff = series.diff().fillna(0) + # fast "no" if the whole series is continuous + if (diff <= 1).all() and (diff >= 0).all(): + return False + + # filter to ignore singleton blips + # e.g. frame_num breaks in one buffer, + # find numbers that don't return to the prior number or number + 1 in the subsequent rows + blips = np.logical_and( + ~diff.between(0, 1), + np.logical_or(diff == diff.shift(-1) * -1, diff == (diff.shift(-1) - 1) * -1), + ) + + # now check if there are any longer lasting discontinuities + diff = series[~blips].diff().dropna() + return bool((~diff.between(0, 1)).any()) diff --git a/mio/process/video.py b/mio/process/video.py index 4ff8d2c7..7a843719 100644 --- a/mio/process/video.py +++ b/mio/process/video.py @@ -8,7 +8,7 @@ import cv2 import numpy as np import pandas as pd -from tqdm import tqdm +from tqdm import tqdm, trange from mio import init_logger from mio.io import VideoReader, VideoWriter @@ -21,7 +21,11 @@ NoisePatchConfig, ) from mio.plots.video import VideoPlotter -from mio.process.frame_helper import FrequencyMaskHelper, InvalidFrameDetector +from mio.process.frame_helper import ( + FrequencyMaskHelper, + InvalidFrameDetector, + make_detectors, +) from mio.process.zstack_helper import ZStackHelper logger = init_logger("video") @@ -131,16 +135,10 @@ def __init__( self.noise_detect_helper = InvalidFrameDetector(noise_patch_config=noise_patch_config) self.noise_patchs: list[np.ndarray] = [] self.noisy_frames: list[np.ndarray] = [] - self.diff_frames: list[np.ndarray] = [] self.dropped_frame_indices: list[int] = [] self.output_enable: bool = noise_patch_config.output_result - if "mean_error" in noise_patch_config.method: - logger.warning( - "The mean_error method is unstable and not fully tested yet." " Use with caution." - ) - def process_frame(self, input_frame: np.ndarray, index: int) -> np.ndarray | None: """ Process a single frame. @@ -185,15 +183,6 @@ def noise_patch_named_video(self) -> NamedVideo: """ return NamedVideo(name="patched_area", video=self.noise_patchs) - @property - def diff_frames_named_video(self) -> NamedVideo: - """ - Get the NamedFrame object for the difference frames. - """ - if not hasattr(self.noise_patch_config, "diff_multiply"): - diff_multiply = 1 - return NamedVideo(name=f"diff_{diff_multiply}x", video=self.diff_frames) - @property def noisy_frames_named_video(self) -> NamedVideo: """ @@ -217,18 +206,6 @@ def export_noise_patch(self) -> None: else: logger.info(f"{self.name} noise patch output disabled.") - def export_diff_frames(self) -> None: - """ - Export the difference frames to a file. - """ - if self.noise_patch_config.output_diff: - logger.info(f"Exporting {self.name} difference frames to {self.output_dir}") - self.diff_frames_named_video.export( - output_path=self.output_dir / f"{self.name}", fps=20, suffix=True, force=self.force - ) - else: - logger.info(f"{self.name} difference frames output disabled.") - def export_noisy_video(self) -> None: """ Export the noisy frames to a file. @@ -250,7 +227,6 @@ def batch_export_videos(self) -> None: """ self.export_output_video() self.export_noise_patch() - self.export_diff_frames() self.export_noisy_video() @@ -583,7 +559,6 @@ def denoise( noise_patch_processor.output_dir = debug_dir noise_patch_processor.export_noise_patch() - noise_patch_processor.export_diff_frames() noise_patch_processor.export_noisy_video() freq_mask_processor.batch_export_videos() @@ -635,6 +610,44 @@ def denoise( return output_video_path +def score_noise( + recording: Recording, config: NoisePatchConfig, progress: bool = False +) -> pd.DataFrame: + """ + Score framewise noise from a recording, + yielding a dataframe with columns for each kind of noise + + - reconstructed_frame_index: the index of the frame in the video + - gradient: the number of pixels that are part of noise patches, + as determined by the second row-wise derivative being above the configured threshold + - black_area: the number of pixels in contiguous black regions (and thus missing) + """ + + records = [] + detectors = make_detectors(config) + n_frames = recording.video.n_frames + iterator = trange(n_frames) if progress else range(n_frames) + + try: + for idx in iterator: + record = {"reconstructed_frame_index": idx} + frame = recording.video[idx] + # FIXME: video proxy should return grayscale as grayscale + # https://github.com/miniscope/mio/issues/175 + if len(frame.shape) == 3 and frame.shape[-1] == 3: + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) + for method, detector in detectors.items(): + _, noise_mask = detector.find_invalid_area(frame) + record[method] = np.count_nonzero(noise_mask) + records.append(record) + + finally: + if progress: + iterator.close() + + return pd.DataFrame(records) + + def trim( video_path: Path, output_path: Path | None = None, diff --git a/mio/utils.py b/mio/utils.py index bcb25950..457ec51e 100644 --- a/mio/utils.py +++ b/mio/utils.py @@ -49,6 +49,8 @@ def hash_video( Returns: str """ + if not Path(path).exists(): + raise FileNotFoundError("No such video exists!") h = hashlib.new(method) vid = cv2.VideoCapture(str(path)) diff --git a/pyproject.toml b/pyproject.toml index c10fc52e..1f58dfb2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,7 +100,7 @@ format.composite = [ "black mio", "ruff check --fix", ] -docs-prod = "sphinx-build -M html ./docs ./docs/_build -W -P -E -a" +docs-prod = "python -m sphinx -M html ./docs ./docs/_build -W -E -a --keep-going" [tool.pdm.build] includes = ["mio"] diff --git a/tests/data/config/denoise_noise_detection_test.yml b/tests/data/config/denoise_noise_detection_test.yml index 886bf644..bc10f898 100644 --- a/tests/data/config/denoise_noise_detection_test.yml +++ b/tests/data/config/denoise_noise_detection_test.yml @@ -4,12 +4,6 @@ mio_version: 0.6.1 noise_patch: enable: true method: [gradient, black_area] - mean_error_config: - threshold: 40 - device_config_id: wireless-200px - buffer_split: 8 - comparison_unit: 1000 - diff_multiply: 1 gradient_config: threshold: 20 black_area_config: diff --git a/tests/test_cli_process.py b/tests/test_cli_process.py index 36a1626d..eaacef19 100644 --- a/tests/test_cli_process.py +++ b/tests/test_cli_process.py @@ -15,10 +15,10 @@ STITCH_DATA_DIR = Path(__file__).parent / "data" / "stitch" -EXPECTED_STITCHED_VIDEO_HASH = "c8cdf3149f812ae25e6f3f1a876249e4ce118e9a53aa1805e48b995b01f07a91" +EXPECTED_STITCHED_VIDEO_HASH = "df937c8651cf142b4d8e2a75140729dcacdc1151ebc3767b48d0ca71578007ff" EXPECTED_CROP_VIDEO_HASH = "432642b1528fcd9ad553cfb3cc3862bef931301bd11d44dc3c2372fc379fa629" EXPECTED_STITCHED_TRIMMED_VIDEO_HASH = ( - "2c62b65ddd537e94e7d3f29e7c46523357d70aefed02d46baa9726ee57798af9" + "d7b6858c85e13da69921593570a8eddff3716d0be348219d197f216afbe3867a" ) EXPECTED_REMOVE_FRAMES_HASH = "b76b80f45316bad0a808802b8f5c0d65b99f6f59bc6422b84c1c2a7026ca4b15" diff --git a/tests/test_process/test_frame_helper.py b/tests/test_process/test_frame_helper.py index 41b7ad55..131bfc0e 100644 --- a/tests/test_process/test_frame_helper.py +++ b/tests/test_process/test_frame_helper.py @@ -7,8 +7,8 @@ from pprint import pformat from pydantic import BaseModel -from mio.models.process import DenoiseConfig, NoisePatchConfig -from mio.process.frame_helper import InvalidFrameDetector +from mio.models.process import BlackAreaDetectorConfig, DenoiseConfig, NoisePatchConfig +from mio.process.frame_helper import BlackAreaDetector, InvalidFrameDetector from ..conftest import DATA_DIR @@ -50,7 +50,6 @@ class NoiseGroundTruth(BaseModel): [ (["gradient"], GroundTruthCategory.check_pattern), (["black_area"], GroundTruthCategory.blacked_out), - (["mean_error"], GroundTruthCategory.check_pattern), ], ) def test_noisy_frame_detection(video, ground_truth, noise_detection_method, noise_category): @@ -60,14 +59,6 @@ def test_noisy_frame_detection(video, ground_truth, noise_detection_method, nois """ if "gradient" in noise_detection_method: global_config: DenoiseConfig = DenoiseConfig.from_id("denoise_noise_detection_test") - elif "mean_error" in noise_detection_method: - if "extended" in video: - # FIXME: resolve this before merging `feat-preprocess` to `main` - pytest.xfail( - "Bug in comparison to previous frames when first frame is noisy, " - "see https://github.com/Aharoni-Lab/mio/pull/97" - ) - global_config: DenoiseConfig = DenoiseConfig.from_id("denoise_example_mean_error") elif "black_area" in noise_detection_method: global_config: DenoiseConfig = DenoiseConfig.from_id("denoise_noise_detection_test") else: @@ -122,3 +113,30 @@ def test_noisy_frame_detection(video, ground_truth, noise_detection_method, nois ) extra_frames = set(detected_frame_indices) - all_expected assert extra_frames == set(), f"Detected extra, non-noise frames as noisy: {extra_frames}" + + +@pytest.mark.parametrize( + "min_rows,expected_noisy", + [ + (1, True), # default: any flagged row triggers detection + (5, True), # exactly 5 noisy rows meets the threshold + (10, False), # only 5 noisy rows, below threshold of 10 + ], +) +def test_black_area_min_rows(min_rows, expected_noisy): + """min_rows controls how many flagged rows are needed to mark a frame as invalid.""" + # Create a 50x50 frame with 5 rows of consecutive zeros (noisy) and the rest bright + frame = np.ones((50, 50), dtype=np.uint8) * 128 + for row in range(5): + frame[row, :30] = 0 # 30 consecutive black pixels in rows 0-4 + + config = BlackAreaDetectorConfig( + consecutive_threshold=10, + value_threshold=0, + min_rows=min_rows, + ) + detector = BlackAreaDetector(config) + is_noisy, mask = detector.find_invalid_area(frame) + assert ( + is_noisy == expected_noisy + ), f"min_rows={min_rows}: expected noisy={expected_noisy}, got {is_noisy}" diff --git a/tests/test_process/test_stitch.py b/tests/test_process/test_stitch.py index bdc87315..6f5b9690 100644 --- a/tests/test_process/test_stitch.py +++ b/tests/test_process/test_stitch.py @@ -17,7 +17,10 @@ from mio.process.stitch import ( CandidateFrame, StitchRecord, + concat_recordings, stitch, + _align_by_time, + _has_discontinuous_runs, _score_edges, ) from mio.process.video import trim, remove_frames @@ -25,7 +28,7 @@ STITCH_DATA_DIR = Path(__file__).parent.parent / "data" / "stitch" -EXPECTED_STITCHED_VIDEO_HASH = "c8cdf3149f812ae25e6f3f1a876249e4ce118e9a53aa1805e48b995b01f07a91" +EXPECTED_STITCHED_VIDEO_HASH = "df937c8651cf142b4d8e2a75140729dcacdc1151ebc3767b48d0ca71578007ff" EXPECTED_DEBUG_VIDEO_HASH = ( "856e6e5c538532bd0fcfb942616686a5cd262aadb51dd8796adf5de69215c94b", "a69b6cadf4ab1dd8a1097d2c1be298397206db235fd4c5f68febd1700f15a4b6", @@ -119,8 +122,10 @@ def test_score_csv_edge_scoring_tiebreaker(stitch_result: StitchedRecording): df = stitch_result.scores # filter frames where only one video or the other had them df = df[~df["compare_video"].isna()] - # there should be four frames that could be decided on metadata alone - assert len(df[df["selected_edge_score"].isna()]) == 4 + # there should be 7 frames that could be decided on metadata alone + # - 4x on buffer count + # - 3x on black pixels + assert len(df[df["selected_edge_score"].isna()]) == 7 # for all those that had to use edge scores, the selected should be greater or equal edges_scored = df[~df["selected_edge_score"].isna()] for _, row in edges_scored.iterrows(): @@ -287,3 +292,123 @@ def test_remove_frames_invalid(tmp_path): remove_indices=list(range(EXPECTED_VIDEO1_FRAME_COUNT)), output_path=tmp_path / "out.avi", ) + + +def test_concat_recordings(tmp_path, recordings): + """Concatenating two recordings produces contiguous frame indices and correct frame count.""" + + combined_video = tmp_path / "combined.avi" + + combined = concat_recordings( + recordings=list(recordings.values()), + output_video_path=combined_video, + ) + + # Video frame count should be sum of both inputs + expected_frames = sum(r.video.n_frames for r in recordings.values()) + # the Recording class also validates that the metadata has a matching length if present + # so its presence means that the metadata is also matching + assert combined.metadata is not None + + actual_frames = combined.video.n_frames + assert actual_frames == expected_frames + + # CSV should have contiguous reconstructed_frame_index + df = combined.metadata + diffs = df["reconstructed_frame_index"].diff().iloc[1:].to_numpy() + assert (diffs <= 1).all() and (diffs >= 0).all() + + +@pytest.mark.parametrize("flip", [True, False]) +def test_align_by_time(tmp_path, flip): + """Aligning by timestamp finds the inner join of the closest matching timestamps""" + # one normal video with some linspaced times + left_idxes = np.ravel(np.repeat(np.arange(50), 5)) + left_times = np.linspace(0, 1, len(left_idxes)) + left = pd.DataFrame( + {"reconstructed_frame_index": left_idxes, "buffer_recv_unix_time": left_times} + ) + + # one offset video with a blippy frame from a bit flip in the frame_num + right_idxes = np.ravel(np.repeat(np.arange(25), 5)) + right_idxes = np.concat([right_idxes, [25]], axis=0) + right_idxes = np.concat([right_idxes, np.ravel(np.repeat(np.arange(26, 52), 5))], axis=0) + # make same size so sampling rate is the same + right_idxes = right_idxes[: len(left_idxes)] + right_times = np.linspace(0, 1, len(right_idxes)) + 0.1 + right = pd.DataFrame( + {"reconstructed_frame_index": right_idxes, "buffer_recv_unix_time": right_times} + ) + + good, bad = "video1", "video2" + if flip: + good, bad = "video2", "video1" + recordings = [ + Recording.model_construct(name=good, metadata=left), + Recording.model_construct(name=bad, metadata=right), + ] + + aligned = _align_by_time(recordings) + # we should have received 45 frames: 50 frames in the original - 5 frames in 0.1 seconds of lag + assert len(aligned) == 45 + assert np.array_equal(aligned[good], np.arange(5, 50)) + + # we should have dropped frame 25 in the right one + assert 25 not in np.array(aligned[bad]) + assert np.array_equal(aligned[bad], np.concat([np.arange(25), np.arange(26, 46)])) + + +def test_stitch_with_timestamps(stitch_result, tmp_path): + """ + When we scramble the `frame_num`, we can stitch by timestamps. + We should get the same result as if we were able to use frame_num in this case. + """ + # use a temporary version of the recordings because we are going to wreck the metadata + recordings = { + "video1": Recording.from_video(STITCH_DATA_DIR / "video1.avi"), + "video2": Recording.from_video(STITCH_DATA_DIR / "video2.avi"), + } + recordings["video1"].metadata["frame_num"] = np.random.default_rng().integers( + 0, 1000, size=len(recordings["video1"].metadata) + ) + recordings["video2"].metadata["frame_num"] = np.random.default_rng().integers( + 0, 1000, size=len(recordings["video2"].metadata) + ) + + result = stitch(list(recordings.values()), debug_video=True, output_dir=tmp_path) + # we should have an inner join on the frames - so only those without a comparison frame + expected = stitch_result.scores[~stitch_result.scores["compare_video"].isna()] + assert np.array_equal(result.scores["selected_video"], expected["selected_video"]) + + +@pytest.mark.parametrize( + "series,expected", + [ + pytest.param([1, 1, 1, 2, 2, 2, 3, 3, 3], False, id="contiguous-buffers"), + pytest.param([1, 2, 3, 4, 5], False, id="contiguous-frames"), + pytest.param([1, 1, 1, 2, 500, 2, 3, 3, 3], False, id="single-bitflip-same-frame"), + pytest.param([1, 1, 1, 2, 2, 500, 3, 3, 3], False, id="single-bitflip-next-frame"), + pytest.param( + [10, 10, 10, 11, 11, 11, 2, 2, 2, 3, 3, 3], True, id="discontiguous-buffers-lower" + ), + pytest.param( + [10, 10, 10, 11, 11, 11, 20, 20, 20, 21, 21, 21], + True, + id="discontiguous-buffers-higher", + ), + pytest.param([1, 2, 3, 4, 5, 1, 2, 3, 4, 5], True, id="discontiguous-frames-lower"), + pytest.param([1, 2, 3, 4, 5, 10, 11, 12, 13], True, id="discontiguous-frames-higher"), + ], +) +def test_has_discontinuous_runs(series, expected): + """ + We can determine when some timeseries has discontinuous runs, + ignoring when a single value blips incorrectly (like a bit flip in a metadata header). + """ + series = pd.Series(series) + assert _has_discontinuous_runs(series) == expected + + +def test_test_data_is_considered_continuous(recordings): + """Just testing the assumptions of the tests ios all""" + assert not any(_has_discontinuous_runs(r.metadata["frame_num"]) for r in recordings.values()) diff --git a/tests/test_process/test_video.py b/tests/test_process/test_video.py index 6bb70f58..d1d797db 100644 --- a/tests/test_process/test_video.py +++ b/tests/test_process/test_video.py @@ -46,18 +46,6 @@ def test_noise_patch_processor(video_frame, tmp_path): assert processor.output_enable -def test_noise_patch_processor_no_config(random_8bit_video_frame, tmp_path): - denoise_config = DenoiseConfig.from_id("denoise_example") - denoise_config.noise_patch.enable = True - denoise_config.noise_patch.mean_error_config = None - denoise_config.noise_patch.gradient_config = None - denoise_config.noise_patch.black_area_config = None - - # This should raise a ValueError because the necessary configs are not provided - with pytest.raises(ValueError): - NoisePatchProcessor("denoise_example", denoise_config.noise_patch, tmp_path) - - def test_noise_patch_processor_no_methods(random_8bit_video_frame, tmp_path): denoise_config = DenoiseConfig.from_id("denoise_example") denoise_config.noise_patch.enable = True