commit 121bec58b0fb7b52eead028fbc606800382fa9bc Author: SeanBRVFX Date: Mon Mar 2 10:00:09 2026 +0100 Initial ComfyUI-CorridorKey custom node diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f02b17f --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +__pycache__/ +*.py[cod] +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ +.venv/ +dist/ +build/ +*.egg-info/ diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..3bc49b0 --- /dev/null +++ b/LICENSE @@ -0,0 +1,25 @@ +ComfyUI-CorridorKey License Notice + +This custom node is intended to follow the licensing expectations stated by the +upstream CorridorKey repository by nikopueringer: + +https://github.com/nikopueringer/CorridorKey + +Summary of the upstream terms as described in the upstream README: + +1. The upstream project is not presented as a standard standalone SPDX license + file. The README describes the release as effectively a variation of + CC BY-NC-SA 4.0. +2. Repackaging and reselling the tool is not permitted. +3. Offering paid API inference based on the model is not permitted without a + separate agreement from the upstream author. +4. Released variations or improvements should remain free and open source. +5. Future forks or releases should keep the "Corridor Key" name. + +This file is a convenience summary only. The authoritative reference for these +terms is the upstream README section titled: + +"CorridorKey Licensing and Permissions" + +If the upstream author publishes a formal license file or updated terms, those +terms should be reviewed and this file should be updated accordingly. diff --git a/README.md b/README.md new file mode 100644 index 0000000..e5e0388 --- /dev/null +++ b/README.md @@ -0,0 +1,468 @@ +# ComfyUI-CorridorKey + +## 1) Project name + one line summary + +`ComfyUI-CorridorKey` is a ComfyUI custom node package that exposes a native `CorridorKey` node for CorridorKey-style coarse-mask refinement inside ComfyUI. + +## 2) Problem statement and goals (what it does, what it does NOT do) + +`ComfyUI-CorridorKey` adapts the intent of the upstream `CorridorKey` project for ComfyUI users who already build masks inside ComfyUI. Instead of running a separate command-line application and instead of bundling extra mask-generation tools such as `VideoMaMa` or `GVM`, this node accepts: + +- an `IMAGE` input +- a coarse `MASK` input produced by other ComfyUI nodes + +It then performs an edge-aware matte refinement pass and returns: + +- a refined alpha `MASK` +- a selected preview `IMAGE` based on the requested output mode +- dedicated `IMAGE` outputs for foreground, processed, and comp previews + +Goals: + +- Keep the integration native to ComfyUI's standard custom-node format. +- Reuse existing ComfyUI nodes for coarse mask creation. +- Keep runtime dependencies small and compatible with a typical ComfyUI Python environment. +- Provide a deterministic, testable baseline refinement pipeline that is safe to run locally. +- Track upstream `CorridorKey` changes safely and surface newer verified commits without blindly overwriting local custom-node code. + +Non-goals: + +- This project does not reimplement the full upstream research model or ship upstream CorridorKey checkpoints. +- This project does not recreate `VideoMaMa`, `GVM`, or any separate mask-proposal pipeline. +- This project does not download models, call external services, or require network access at runtime. +- This project does not write output files to disk; it only returns ComfyUI node outputs. +- This project does not auto-overwrite its own code from the upstream `CorridorKey` repository, because that repository is not the same codebase as this ComfyUI custom node. + +## 3) Features (current + clearly marked planned) + +Current features: + +- One ComfyUI node: `CorridorKey` +- Accepts batched `IMAGE` tensors and `MASK` tensors +- Broadcasts a single mask across an image batch when needed +- Validates shapes, value ranges, and numeric settings +- Exposes upstream-inspired controls for: + - `Gamma Space` + - `Despill Strength` + - `Auto-Despeckle` + - `Despeckle Size` + - `Refiner Strength` + - output mode selection (`FG`, `Matte`, `Processed`, `Comp`) +- Performs deterministic edge-aware mask smoothing, spill suppression, and comp preview generation +- Returns refined alpha mask plus selected output, foreground, processed, and comp preview images +- On module load, performs a best-effort background check for newer upstream commits that have no failing GitHub check-runs +- Includes pytest coverage for the core processor + +Planned features: + +- Planned: optional backend that can call a true upstream CorridorKey-compatible runtime if the user provides one +- Planned: trimap output for downstream compositing workflows +- Planned: preset profiles tuned for hair, glass, and semi-transparent edges +- Planned: optional preview compositing against solid-color backgrounds + +## 4) Requirements (Python version, OS, dependencies) + +- Python: 3.10 or newer +- OS: Windows, Linux, or macOS where ComfyUI runs +- Runtime dependencies: + - `torch` (normally already present in the ComfyUI Python environment) + - `numpy` + - `Pillow` +- Development dependencies: + - `pytest` + - `ruff` + - `black` + - `mypy` + +## 5) Quick start (venv, install, run) + +This repository is meant to live under ComfyUI's `custom_nodes` directory. + +Recommended ComfyUI installation: + +1. Clone this repository into `ComfyUI/custom_nodes/`. +2. Install the runtime dependencies into the same Python environment ComfyUI uses. +3. Restart ComfyUI. + +Typical install flow: + +```powershell +cd ComfyUI\custom_nodes +git clone ComfyUI-CorridorKey +cd ComfyUI-CorridorKey +python -m pip install -r requirements.txt +``` + +After that, restart ComfyUI and the `CorridorKey` node should appear. + +If you are using ComfyUI Portable on Windows and want to be explicit about the embedded Python: + +```powershell +cd ComfyUI_windows_portable\ComfyUI\custom_nodes +git clone ComfyUI-CorridorKey +..\..\python_embeded\python.exe -m pip install -r .\ComfyUI-CorridorKey\requirements.txt +``` + +Manual install also works: + +1. Copy the `ComfyUI-CorridorKey` folder into `ComfyUI/custom_nodes/`. +2. Run `python -m pip install -r requirements.txt` from inside that folder. +3. Restart ComfyUI. + +Development install: + +```powershell +python -m venv .venv +.venv\Scripts\Activate.ps1 +python -m pip install --upgrade pip +python -m pip install -r requirements.txt +python -m pip install -e .[dev] +``` + +Restart ComfyUI after installation. The node will appear under the `CorridorKey` category. + +For a normal ComfyUI install, `requirements.txt` is the primary runtime dependency file. `pyproject.toml` is kept for editable installs and development tooling. + +## 6) Configuration (env vars, config files, defaults) + +The node does not require config files. + +Runtime configuration is passed through node inputs: + +- `gamma_space`: `sRGB` or `Linear`, matching the upstream input interpretation concept +- `feather_radius`: integer blur radius used to soften the coarse mask before refinement +- `edge_focus`: float multiplier that strengthens image-edge preservation +- `threshold`: float cutoff used to center the matte transition +- `preserve_core`: float band that locks obvious foreground/background values +- `despill_strength`: normalized float from `0.0` to `1.0` that reduces green spill on semi-transparent foreground edges +- `refiner_strength`: float multiplier that strengthens the edge-aware matte correction pass +- `auto_despeckle`: enables a small morphology-based cleanup pass +- `despeckle_size`: integer control for the minimum speckle scale to suppress +- `output_mode`: selects the primary preview output as `FG`, `Matte`, `Processed`, or `Comp` + +Defaults: + +- `gamma_space = sRGB` +- `feather_radius = 4` +- `edge_focus = 1.5` +- `threshold = 0.5` +- `preserve_core = 0.2` +- `despill_strength = 1.0` +- `refiner_strength = 1.0` +- `auto_despeckle = enabled` +- `despeckle_size = 400` +- `output_mode = Processed` + +Implementation note: + +- The upstream app uses model-driven processing. This ComfyUI node mirrors the same user-facing control surface where practical, but the internal math is a deterministic local approximation built around the provided image and coarse mask. +- The upstream CLI uses a numeric prompt for `Auto-Despeckle Size` with a default of `400` pixels. This node exposes the same value as an integer input rather than forcing a fixed dropdown. + +Environment variables for upstream tracking: + +- `CORRIDORKEY_AUTO_CHECK_UPSTREAM` + - Default: `1` + - Set to `0` to disable the background upstream check +- `CORRIDORKEY_UPSTREAM_TIMEOUT_SECONDS` + - Default: `3.0` + - Per-request timeout for GitHub API calls +- `CORRIDORKEY_UPSTREAM_CHECK_DEPTH` + - Default: `15` + - How many recent upstream commits to scan when looking for the newest verified commit + +The currently reviewed upstream head SHA and its last observed check state are pinned in `corridor_key/upstream_sync.py`. + +## 7) Usage (examples, CLI/API examples if relevant) + +Recommended ComfyUI flow: + +1. Load or generate an image. +2. Create a coarse mask using existing ComfyUI nodes such as RMBG, SAM, Segment Anything, Florence-based tools, or manual masking nodes. +3. Feed the `IMAGE` and coarse `MASK` into `CorridorKey`. +4. Choose `FG`, `Matte`, `Processed`, or `Comp` as the node's selected preview output. +5. Use the refined `MASK` for compositing, inpainting, EXR assembly, or export. + +Node inputs: + +- `image`: `IMAGE` +- `mask`: `MASK` +- `gamma_space`: `COMBO` +- `feather_radius`: `INT` +- `edge_focus`: `FLOAT` +- `threshold`: `FLOAT` +- `preserve_core`: `FLOAT` +- `despill_strength`: `FLOAT` +- `refiner_strength`: `FLOAT` +- `auto_despeckle`: `COMBO` +- `despeckle_size`: `INT` +- `output_mode`: `COMBO` + +Node outputs: + +- `refined_mask`: `MASK` +- `selected_output`: `IMAGE` +- `foreground`: `IMAGE` +- `processed`: `IMAGE` +- `comp`: `IMAGE` + +Example workflow idea: + +- `Load Image` -> `RMBG` (or another segmentation node) -> `CorridorKey` -> `Preview Image` +- For EXR-oriented workflows, use the returned `foreground` and `refined_mask` separately and combine or export them with other ComfyUI nodes such as `cocotools_io`, instead of expecting this node to write EXR files directly. + +There is no standalone CLI in this custom node package by design. + +Upstream tracking behavior: + +- On import, the package starts a small background check. +- It queries the upstream `nikopueringer/CorridorKey` GitHub repo. +- It looks for the newest recent commit with at least one successful check-run and no failing check-runs. +- If such a commit is newer than the currently pinned verified upstream baseline, it logs that an upstream update is available for manual review. +- It does not automatically rewrite local files from upstream. + +## 8) Architecture overview + +### High-level diagram in ASCII + +```text +IMAGE + coarse MASK + | + v + input validation + | + v + mask normalization + | + v + optional gamma handling + | + v + image luminance + edge map + | + v + refiner + despeckle + | + v + refined alpha MASK + | + v + despill + output assembly + | | | + v v v + FG Processed Comp +``` + +### Key modules and responsibilities + +- `nodes.py` + - ComfyUI-facing node class + - Defines `INPUT_TYPES`, return types, and node metadata + - Converts node inputs into the internal processor call +- `corridor_key/config.py` + - Dataclass for validated runtime settings +- `corridor_key/tensor_ops.py` + - Shared tensor validation, mask normalization, color helpers, and small image-processing helpers +- `corridor_key/processor.py` + - Pure processing pipeline for edge-aware matte refinement, despill, despeckle, and output assembly +- `corridor_key/upstream_sync.py` + - Best-effort GitHub API integration for verified-upstream commit checks +- `tests/test_processor.py` + - Unit tests for shape handling, validation, and output behavior + +### Data flow / request flow + +1. ComfyUI invokes `CorridorKey.run`. +2. The node builds a `CorridorKeySettings` object from numeric inputs. +3. The processor validates the image and mask tensors. +4. A normalized mask batch is created and aligned to the image batch. +5. The processor optionally adapts handling for `Gamma Space`. +6. The processor derives luminance, computes an edge map, and applies the refiner pass. +7. Optional despeckle cleanup is applied when enabled. +8. The processor builds foreground, processed, and comp previews, including optional despill. +9. The processor returns: + - refined mask + - selected output + - foreground + - processed + - comp +10. ComfyUI receives the outputs without any file-system side effects. +11. In parallel, an optional background task may query GitHub and log if a newer verified upstream commit is available. + +## 9) Repository structure (tree) + +```text +ComfyUI-CorridorKey/ +|-- README.md +|-- LICENSE +|-- requirements.txt +|-- pyproject.toml +|-- __init__.py +|-- nodes.py +|-- corridor_key/ +| |-- __init__.py +| |-- config.py +| |-- tensor_ops.py +| |-- processor.py +| `-- upstream_sync.py +`-- tests/ + `-- test_processor.py +``` + +## 10) Development + +### local dev loop + +```powershell +python -m pip install -e .[dev] +pytest -q +ruff check . +black --check . +mypy corridor_key +``` + +When changing behavior: + +1. Update `README.md` first if the behavior, architecture, setup, or workflow changes. +2. Keep the implementation aligned with the documented repository structure. +3. Add or update pytest coverage for the changed logic. +4. Restart ComfyUI and re-test the node in a simple image-plus-mask workflow. + +### testing strategy (pytest) + +- Test pure processing code in `corridor_key/processor.py` +- Avoid ComfyUI runtime dependencies in the test suite where possible +- Cover: + - mask broadcasting + - parameter validation + - output shapes + - value range clamping + - output mode selection + - despeckle and despill edge cases + - verified-upstream selection policy for mocked GitHub metadata + +### lint/format rules (ruff + black or ruff format) + +- `ruff check .` for linting +- `black .` for formatting +- Keep imports explicit and functions small +- Prefer readable tensor code over compact but opaque expressions + +### type checking (mypy or pyright if useful) + +- `mypy corridor_key` +- Use standard type hints for public helpers and processor methods +- Keep the ComfyUI adapter lightweight and typed where practical + +## 11) Error handling and logging (logging module, structured messages) + +- Internal validation raises `ValueError` with specific, user-readable messages +- The ComfyUI node catches nothing silently; invalid inputs should fail clearly +- Logging is intentionally minimal: + - a module-level logger may emit debug-level messages in the processor + - a background upstream checker may emit info-level status lines + - no print-based progress logging + - GitHub API checks use explicit short timeouts and fail closed + +## 12) Security notes (inputs, secrets, safe defaults) + +- No `eval`, `exec`, or unsafe deserialization +- No subprocess execution +- No model downloads or arbitrary network access during image processing +- No secrets are required +- All user-provided numeric inputs are validated and clamped to safe ranges +- File paths are not accepted as node inputs, so the node does not read or write arbitrary files +- The processor only operates on in-memory tensors returned by ComfyUI +- The optional upstream checker uses only GitHub's API with a short timeout, does not require tokens, and never executes remote code + +## 13) Deployment (if applicable) + +Deployment is simply installation as a ComfyUI custom node: + +1. Place this folder under `ComfyUI/custom_nodes/`. +2. Install the package into the same Python environment used by ComfyUI. +3. Restart ComfyUI. + +There is no separate server, container, or packaged service in this repository. + +## 14) Troubleshooting and FAQ + +Q: Why does this node require a mask input? + +A: This package intentionally reuses existing ComfyUI segmentation and masking nodes instead of duplicating them. + +Q: Why does it not match the full upstream CorridorKey model exactly? + +A: This implementation provides a native, deterministic refinement pass inspired by the same workflow goal, without bundling upstream research weights or optional external pipelines. + +Q: My mask looks too soft. + +A: Lower `feather_radius`, lower `preserve_core`, raise `edge_focus`, or raise `refiner_strength`. + +Q: My mask looks too harsh. + +A: Raise `feather_radius` slightly or lower `threshold`. + +Q: Why is `Processed` not writing an EXR with alpha by itself? + +A: ComfyUI nodes exchange in-memory tensors. This node returns separate `foreground` and `refined_mask` outputs so you can assemble or export EXR data elsewhere in the workflow. + +Q: Why is the default `despeckle_size` so large? + +A: The upstream `CorridorKey` CLI defaults `Auto-Despeckle Size` to `400` pixels. On small images, you may need to lower that value or disable auto-despeckle. + +Q: Does this node auto-update itself from GitHub? + +A: It automatically checks for newer upstream commits that appear verified, but it does not self-modify local code. The upstream repo is a different project, so updates should still be reviewed and ported intentionally. + +## 15) Roadmap + +- Add an optional pluggable backend for a true upstream-compatible inference engine +- Add trimap and confidence outputs +- Add more regression tests around semi-transparent regions +- Add a true RGBA or EXR-friendly packed output path if the surrounding ComfyUI ecosystem provides a stable tensor contract for it +- Add an example ComfyUI workflow JSON once the node behavior stabilizes + +## 16) Contributing (branching, commit style, PR checklist) + +Branching: + +- Use short-lived feature branches from `main` +- Keep one focused change per branch + +Commit style: + +- Prefer imperative commit messages +- Example: `Add edge-aware CorridorKey refinement node` + +PR checklist: + +- Update `README.md` first when behavior or structure changes +- Run `pytest -q` +- Run `ruff check .` +- Run `black --check .` +- Run `mypy corridor_key` +- Confirm the node loads in ComfyUI +- Keep the public node contract backward compatible unless documented + +## 17) License placeholder + +This package is based on the workflow and licensing expectations stated by the upstream `nikopueringer/CorridorKey` repository. + +Upstream license status: + +- The upstream repository does not present a standard standalone SPDX license file in the root. +- Its README includes a custom license statement described by the author as effectively a variation of `CC BY-NC-SA 4.0`. +- The upstream README also states additional restrictions, including: + - no repackaging and resale of the tool + - no paid API inference service using the model without a separate agreement + - released variations or improvements should remain free and open source + - future forks or releases should keep the `Corridor Key` name + +For this custom node, treat the upstream repository's licensing terms as the governing reference unless the repository owner explicitly relicenses this work. + +This repository also includes a local `LICENSE` file that summarizes the same upstream-derived terms for convenience. + +Source reference: + +- Upstream repository: `https://github.com/nikopueringer/CorridorKey` +- Upstream licensing section: `README.md`, section `CorridorKey Licensing and Permissions` diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..c5e6f1e --- /dev/null +++ b/__init__.py @@ -0,0 +1,33 @@ +""" +@author: local +@title: ComfyUI-CorridorKey +@nickname: CorridorKey +@description: ComfyUI node for CorridorKey-style coarse mask refinement. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +try: + from .corridor_key import schedule_upstream_check + from .nodes import CorridorKey +except ImportError: + package_dir = Path(__file__).resolve().parent + if str(package_dir) not in sys.path: + sys.path.insert(0, str(package_dir)) + from corridor_key import schedule_upstream_check + from nodes import CorridorKey + +NODE_CLASS_MAPPINGS = { + "CorridorKey": CorridorKey, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "CorridorKey": "CorridorKey", +} + +schedule_upstream_check() + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/corridor_key/__init__.py b/corridor_key/__init__.py new file mode 100644 index 0000000..d4aed9d --- /dev/null +++ b/corridor_key/__init__.py @@ -0,0 +1,25 @@ +from .config import CorridorKeySettings +from .processor import CorridorKeyProcessor +from .upstream_sync import ( + SYNCED_UPSTREAM_HEAD_CHECK_CONCLUSIONS, + SYNCED_UPSTREAM_HEAD_DATE, + SYNCED_UPSTREAM_HEAD_MESSAGE, + SYNCED_UPSTREAM_HEAD_SHA, + UpstreamCommitRecord, + is_verified_check_conclusions, + schedule_upstream_check, + select_latest_verified_commit, +) + +__all__ = [ + "CorridorKeyProcessor", + "CorridorKeySettings", + "SYNCED_UPSTREAM_HEAD_CHECK_CONCLUSIONS", + "SYNCED_UPSTREAM_HEAD_DATE", + "SYNCED_UPSTREAM_HEAD_MESSAGE", + "SYNCED_UPSTREAM_HEAD_SHA", + "UpstreamCommitRecord", + "is_verified_check_conclusions", + "schedule_upstream_check", + "select_latest_verified_commit", +] diff --git a/corridor_key/config.py b/corridor_key/config.py new file mode 100644 index 0000000..372f93e --- /dev/null +++ b/corridor_key/config.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class CorridorKeySettings: + gamma_space: str = "sRGB" + feather_radius: int = 4 + edge_focus: float = 1.5 + threshold: float = 0.5 + preserve_core: float = 0.2 + despill_strength: float = 1.0 + refiner_strength: float = 1.0 + auto_despeckle: str = "On" + despeckle_size: int = 400 + output_mode: str = "Processed" + + def __post_init__(self) -> None: + if self.gamma_space not in {"sRGB", "Linear"}: + raise ValueError("gamma_space must be 'sRGB' or 'Linear'.") + if not 0 <= self.feather_radius <= 64: + raise ValueError("feather_radius must be between 0 and 64.") + if not 0.0 <= self.edge_focus <= 4.0: + raise ValueError("edge_focus must be between 0.0 and 4.0.") + if not 0.0 <= self.threshold <= 1.0: + raise ValueError("threshold must be between 0.0 and 1.0.") + if not 0.0 <= self.preserve_core <= 0.45: + raise ValueError("preserve_core must be between 0.0 and 0.45.") + if not 0.0 <= self.despill_strength <= 1.0: + raise ValueError("despill_strength must be between 0.0 and 1.0.") + if not 0.1 <= self.refiner_strength <= 4.0: + raise ValueError("refiner_strength must be between 0.1 and 4.0.") + if self.auto_despeckle not in {"Off", "On"}: + raise ValueError("auto_despeckle must be 'Off' or 'On'.") + if not 1 <= self.despeckle_size <= 4096: + raise ValueError("despeckle_size must be between 1 and 4096.") + if self.output_mode not in {"FG", "Matte", "Processed", "Comp"}: + raise ValueError("output_mode must be one of: FG, Matte, Processed, Comp.") + + @property + def transition_gain(self) -> float: + return 8.0 + (self.edge_focus * 16.0) + + @property + def despill_factor(self) -> float: + return self.despill_strength + + @property + def despeckle_enabled(self) -> bool: + return self.auto_despeckle == "On" diff --git a/corridor_key/processor.py b/corridor_key/processor.py new file mode 100644 index 0000000..7c88876 --- /dev/null +++ b/corridor_key/processor.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import logging + +import torch + +from .config import CorridorKeySettings +from .tensor_ops import ( + build_checkerboard, + clamp_unit_interval, + compute_edge_map, + convert_gamma_space, + despeckle_binary_mask, + ensure_image_tensor, + ensure_mask_batch, + gaussian_blur_map, + mask_to_image, +) + +LOGGER = logging.getLogger(__name__) + + +class CorridorKeyProcessor: + # Mini title: lightweight green-spill suppression for edge pixels. + def _apply_despill(self, image: torch.Tensor, alpha: torch.Tensor, strength: float) -> torch.Tensor: + if strength <= 0.0: + return image + + edge_weight = 1.0 - (alpha.mul(2.0).sub(1.0).abs()) + rb_max = torch.maximum(image[..., 0], image[..., 2]) + spill = torch.relu(image[..., 1] - rb_max) + reduction = spill * strength * (0.25 + (0.75 * edge_weight)) + + red = clamp_unit_interval(image[..., 0] + (reduction * 0.10)) + green = clamp_unit_interval(image[..., 1] - reduction) + blue = clamp_unit_interval(image[..., 2] + (reduction * 0.10)) + return torch.stack((red, green, blue), dim=-1) + + # Mini title: choose the user-requested primary preview. + def _select_output( + self, + output_mode: str, + alpha: torch.Tensor, + foreground: torch.Tensor, + processed: torch.Tensor, + comp: torch.Tensor, + ) -> torch.Tensor: + if output_mode == "FG": + return foreground + if output_mode == "Matte": + return mask_to_image(alpha) + if output_mode == "Comp": + return comp + return processed + + def refine( + self, + image: torch.Tensor, + mask: torch.Tensor, + settings: CorridorKeySettings, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if not isinstance(settings, CorridorKeySettings): + raise ValueError("settings must be a CorridorKeySettings instance.") + + image_batch = ensure_image_tensor(image) + working_image = convert_gamma_space(image_batch, settings.gamma_space) + mask_batch = ensure_mask_batch( + mask=mask, + batch_size=working_image.shape[0], + height=working_image.shape[1], + width=working_image.shape[2], + ) + + # Mini title: compute the soft starting matte. + softened_mask = gaussian_blur_map(mask_batch, settings.feather_radius) + + # Mini title: derive edge guidance from image luminance. + edge_map = compute_edge_map(working_image) + edge_gain = 1.0 + (edge_map * settings.edge_focus) + + centered_mask = ( + (softened_mask - settings.threshold) + * settings.transition_gain + * settings.refiner_strength + ) + alpha = torch.sigmoid(centered_mask * edge_gain) + + low_lock = max(0.0, settings.threshold - settings.preserve_core) + high_lock = min(1.0, settings.threshold + settings.preserve_core) + alpha = torch.where(mask_batch <= low_lock, torch.zeros_like(alpha), alpha) + alpha = torch.where(mask_batch >= high_lock, torch.ones_like(alpha), alpha) + + # Mini title: suppress tiny disconnected islands when requested. + if settings.despeckle_enabled: + support_mask = despeckle_binary_mask( + (alpha >= settings.threshold).to(dtype=alpha.dtype), + settings.despeckle_size, + ) + alpha = torch.where(support_mask > 0.0, alpha, torch.zeros_like(alpha)) + + alpha = clamp_unit_interval(alpha) + + # Mini title: build the foreground-style outputs. + despilled_image = self._apply_despill(working_image, alpha, settings.despill_factor) + visible_support = (alpha > 0.01).to(dtype=despilled_image.dtype).unsqueeze(-1) + foreground = clamp_unit_interval(despilled_image * visible_support) + + alpha_rgb = alpha.unsqueeze(-1) + processed = clamp_unit_interval(foreground * alpha_rgb) + checkerboard = build_checkerboard( + batch_size=working_image.shape[0], + height=working_image.shape[1], + width=working_image.shape[2], + device=working_image.device, + dtype=working_image.dtype, + ) + comp = clamp_unit_interval(processed + (checkerboard * (1.0 - alpha_rgb))) + selected_output = self._select_output( + output_mode=settings.output_mode, + alpha=alpha, + foreground=foreground, + processed=processed, + comp=comp, + ) + + LOGGER.debug( + "CorridorKey refinement complete for batch=%s height=%s width=%s", + working_image.shape[0], + working_image.shape[1], + working_image.shape[2], + ) + + return alpha, selected_output, foreground, processed, comp diff --git a/corridor_key/tensor_ops.py b/corridor_key/tensor_ops.py new file mode 100644 index 0000000..1b99901 --- /dev/null +++ b/corridor_key/tensor_ops.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +from collections import deque + +import numpy as np +import torch +import torch.nn.functional as F + + +def clamp_unit_interval(tensor: torch.Tensor) -> torch.Tensor: + return tensor.to(dtype=torch.float32).clamp(0.0, 1.0) + + +def convert_gamma_space(image: torch.Tensor, gamma_space: str) -> torch.Tensor: + normalized = clamp_unit_interval(image) + if gamma_space == "Linear": + return normalized.clamp_min(0.0).pow(1.0 / 2.2) + return normalized + + +def ensure_image_tensor(image: torch.Tensor) -> torch.Tensor: + if not isinstance(image, torch.Tensor): + raise ValueError("image must be a torch.Tensor.") + if image.ndim != 4 or image.shape[-1] != 3: + raise ValueError("image must have shape [batch, height, width, 3].") + if image.shape[0] < 1 or image.shape[1] < 1 or image.shape[2] < 1: + raise ValueError("image must contain at least one pixel.") + return clamp_unit_interval(image) + + +def ensure_mask_batch( + mask: torch.Tensor, + batch_size: int, + height: int, + width: int, +) -> torch.Tensor: + if not isinstance(mask, torch.Tensor): + raise ValueError("mask must be a torch.Tensor.") + + normalized = mask.to(dtype=torch.float32) + + if normalized.ndim == 2: + normalized = normalized.unsqueeze(0) + elif normalized.ndim == 3: + pass + elif normalized.ndim == 4 and normalized.shape[1] == 1: + normalized = normalized.squeeze(1) + elif normalized.ndim == 4 and normalized.shape[-1] == 1: + normalized = normalized.squeeze(-1) + else: + raise ValueError("mask must have shape [height, width], [batch, height, width], or a single-channel 4D shape.") + + if normalized.ndim != 3: + raise ValueError("mask normalization failed to produce a [batch, height, width] tensor.") + if normalized.shape[-2:] != (height, width): + raise ValueError("mask height and width must match the image tensor.") + + if normalized.shape[0] == 1 and batch_size > 1: + normalized = normalized.repeat(batch_size, 1, 1) + elif normalized.shape[0] != batch_size: + raise ValueError("mask batch dimension must match the image batch or be a single mask.") + + return clamp_unit_interval(normalized) + + +def gaussian_blur_map(mask_batch: torch.Tensor, radius: int) -> torch.Tensor: + if radius <= 0: + return clamp_unit_interval(mask_batch) + + sigma = max(radius / 2.0, 0.5) + positions = torch.arange(-radius, radius + 1, device=mask_batch.device, dtype=mask_batch.dtype) + kernel_1d = torch.exp(-(positions**2) / (2 * sigma * sigma)) + kernel_1d = kernel_1d / kernel_1d.sum() + + kernel_x = kernel_1d.view(1, 1, 1, -1) + kernel_y = kernel_1d.view(1, 1, -1, 1) + + working = mask_batch.unsqueeze(1) + working = F.pad(working, (radius, radius, 0, 0), mode="replicate") + working = F.conv2d(working, kernel_x) + working = F.pad(working, (0, 0, radius, radius), mode="replicate") + working = F.conv2d(working, kernel_y) + return clamp_unit_interval(working.squeeze(1)) + + +def despeckle_binary_mask(mask_batch: torch.Tensor, minimum_size: int) -> torch.Tensor: + if minimum_size <= 1: + return clamp_unit_interval(mask_batch) + + binary = (clamp_unit_interval(mask_batch) >= 0.5).to(dtype=torch.uint8) + output = binary.clone() + cpu_masks = binary.detach().cpu().numpy() + + for batch_index, cpu_mask in enumerate(cpu_masks): + keep_mask = np.zeros_like(cpu_mask, dtype=np.uint8) + visited = np.zeros_like(cpu_mask, dtype=bool) + height, width = cpu_mask.shape + + for y_coord in range(height): + for x_coord in range(width): + if cpu_mask[y_coord, x_coord] == 0 or visited[y_coord, x_coord]: + continue + + queue: deque[tuple[int, int]] = deque([(y_coord, x_coord)]) + component: list[tuple[int, int]] = [] + visited[y_coord, x_coord] = True + + while queue: + current_y, current_x = queue.popleft() + component.append((current_y, current_x)) + + for next_y, next_x in ( + (current_y - 1, current_x), + (current_y + 1, current_x), + (current_y, current_x - 1), + (current_y, current_x + 1), + ): + if next_y < 0 or next_y >= height or next_x < 0 or next_x >= width: + continue + if visited[next_y, next_x] or cpu_mask[next_y, next_x] == 0: + continue + visited[next_y, next_x] = True + queue.append((next_y, next_x)) + + if len(component) >= minimum_size: + for keep_y, keep_x in component: + keep_mask[keep_y, keep_x] = 1 + + output[batch_index] = torch.from_numpy(keep_mask).to(device=mask_batch.device) + + return output.to(dtype=mask_batch.dtype) + + +def compute_edge_map(image: torch.Tensor) -> torch.Tensor: + luminance = ( + (image[..., 0] * 0.299) + + (image[..., 1] * 0.587) + + (image[..., 2] * 0.114) + ).unsqueeze(1) + + sobel_x = image.new_tensor( + [[[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]]] + ).unsqueeze(0) + sobel_y = image.new_tensor( + [[[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0]]] + ).unsqueeze(0) + + padded = F.pad(luminance, (1, 1, 1, 1), mode="replicate") + grad_x = F.conv2d(padded, sobel_x) + grad_y = F.conv2d(padded, sobel_y) + magnitude = torch.sqrt((grad_x.square() + grad_y.square()).clamp_min(1e-12)).squeeze(1) + + batch_max = magnitude.flatten(1).max(dim=1).values.clamp_min(1e-6).view(-1, 1, 1) + return clamp_unit_interval(magnitude / batch_max) + + +def mask_to_image(mask_batch: torch.Tensor) -> torch.Tensor: + grayscale = clamp_unit_interval(mask_batch).unsqueeze(-1) + return grayscale.repeat(1, 1, 1, 3) + + +def build_checkerboard( + batch_size: int, + height: int, + width: int, + device: torch.device, + dtype: torch.dtype, + tile_size: int = 8, +) -> torch.Tensor: + y_index = torch.arange(height, device=device) + x_index = torch.arange(width, device=device) + grid_y = (y_index[:, None] // tile_size) % 2 + grid_x = (x_index[None, :] // tile_size) % 2 + board = ((grid_x + grid_y) % 2).to(dtype=dtype) + board = (board * 0.30) + 0.35 + board = board.unsqueeze(0).unsqueeze(-1).repeat(batch_size, 1, 1, 3) + return clamp_unit_interval(board) diff --git a/corridor_key/upstream_sync.py b/corridor_key/upstream_sync.py new file mode 100644 index 0000000..232c63f --- /dev/null +++ b/corridor_key/upstream_sync.py @@ -0,0 +1,197 @@ +from __future__ import annotations + +import json +import logging +import os +import threading +from dataclasses import dataclass +from typing import Any, Sequence +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + +LOGGER = logging.getLogger(__name__) + +UPSTREAM_API_ROOT = "https://api.github.com/repos/nikopueringer/CorridorKey" +SYNCED_UPSTREAM_HEAD_SHA = "38248c6dd1785b8d7e51ef257fd0889143ca58af" +SYNCED_UPSTREAM_HEAD_DATE = "2026-03-01T22:23:57Z" +SYNCED_UPSTREAM_HEAD_MESSAGE = ( + "Merge pull request #34 from taylorOntologize/tgw/coverage-review" +) +SYNCED_UPSTREAM_HEAD_CHECK_CONCLUSIONS = ("failure", "success", "cancelled") + +_BLOCKING_CHECK_CONCLUSIONS = { + "action_required", + "failure", + "startup_failure", + "timed_out", +} +_CHECK_THREAD_STARTED = False +_CHECK_THREAD_LOCK = threading.Lock() + + +@dataclass(frozen=True, slots=True) +class UpstreamCommitRecord: + sha: str + date: str + message: str + conclusions: tuple[str, ...] + + +def _parse_bool_env(name: str, default: bool) -> bool: + raw_value = os.getenv(name) + if raw_value is None: + return default + return raw_value.strip().lower() not in {"0", "false", "no", "off"} + + +def _parse_float_env(name: str, default: float, minimum: float, maximum: float) -> float: + raw_value = os.getenv(name) + if raw_value is None: + return default + try: + parsed = float(raw_value) + except ValueError: + return default + return max(minimum, min(maximum, parsed)) + + +def _parse_int_env(name: str, default: int, minimum: int, maximum: int) -> int: + raw_value = os.getenv(name) + if raw_value is None: + return default + try: + parsed = int(raw_value) + except ValueError: + return default + return max(minimum, min(maximum, parsed)) + + +def is_verified_check_conclusions(conclusions: Sequence[str]) -> bool: + normalized = tuple( + str(conclusion).strip().lower() + for conclusion in conclusions + if str(conclusion).strip() + ) + if not normalized: + return False + has_success = any(conclusion == "success" for conclusion in normalized) + has_blocker = any(conclusion in _BLOCKING_CHECK_CONCLUSIONS for conclusion in normalized) + return has_success and not has_blocker + + +def _extract_commit_record(commit_payload: dict[str, Any], check_payload: dict[str, Any]) -> UpstreamCommitRecord: + message = str( + commit_payload.get("commit", {}) + .get("message", "") + ).splitlines()[0] + conclusions = tuple( + str(check_run.get("conclusion", "")).strip().lower() + for check_run in check_payload.get("check_runs", []) + ) + return UpstreamCommitRecord( + sha=str(commit_payload.get("sha", "")).strip(), + date=str(commit_payload.get("commit", {}).get("author", {}).get("date", "")).strip(), + message=message, + conclusions=conclusions, + ) + + +def select_latest_verified_commit( + commit_payloads: Sequence[dict[str, Any]], + check_payloads_by_sha: dict[str, dict[str, Any]], +) -> UpstreamCommitRecord | None: + for commit_payload in commit_payloads: + sha = str(commit_payload.get("sha", "")).strip() + if not sha: + continue + check_payload = check_payloads_by_sha.get(sha) + if check_payload is None: + continue + record = _extract_commit_record(commit_payload, check_payload) + if is_verified_check_conclusions(record.conclusions): + return record + return None + + +def _fetch_json(url: str, timeout_seconds: float) -> Any: + request = Request( + url, + headers={ + "Accept": "application/vnd.github+json", + "User-Agent": "ComfyUI-CorridorKey", + }, + method="GET", + ) + with urlopen(request, timeout=timeout_seconds) as response: + return json.loads(response.read().decode("utf-8")) + + +def fetch_latest_verified_commit(timeout_seconds: float, depth: int) -> UpstreamCommitRecord | None: + commits_url = f"{UPSTREAM_API_ROOT}/commits?sha=main&per_page={depth}" + commit_payloads = _fetch_json(commits_url, timeout_seconds) + if not isinstance(commit_payloads, list): + raise ValueError("GitHub commits API returned an unexpected payload.") + + check_payloads_by_sha: dict[str, dict[str, Any]] = {} + for commit_payload in commit_payloads: + sha = str(commit_payload.get("sha", "")).strip() + if not sha: + continue + check_url = f"{UPSTREAM_API_ROOT}/commits/{sha}/check-runs" + check_payload = _fetch_json(check_url, timeout_seconds) + if isinstance(check_payload, dict): + check_payloads_by_sha[sha] = check_payload + + return select_latest_verified_commit(commit_payloads, check_payloads_by_sha) + + +def _run_upstream_check() -> None: + timeout_seconds = _parse_float_env("CORRIDORKEY_UPSTREAM_TIMEOUT_SECONDS", 3.0, 0.5, 10.0) + depth = _parse_int_env("CORRIDORKEY_UPSTREAM_CHECK_DEPTH", 15, 1, 30) + + try: + latest_verified = fetch_latest_verified_commit( + timeout_seconds=timeout_seconds, + depth=depth, + ) + except (HTTPError, URLError, OSError, TimeoutError, ValueError) as exc: + LOGGER.debug("CorridorKey upstream check skipped: %s", exc) + return + + if latest_verified is None: + LOGGER.debug( + "CorridorKey upstream check found no recent commit without failing checks." + ) + return + + if latest_verified.sha == SYNCED_UPSTREAM_HEAD_SHA: + LOGGER.debug( + "CorridorKey upstream head %s is already the current reviewed baseline.", + latest_verified.sha[:12], + ) + return + + LOGGER.info( + "CorridorKey upstream update available: %s (%s). " + "This node does not auto-overwrite itself; review and port changes manually.", + latest_verified.sha[:12], + latest_verified.date, + ) + + +def schedule_upstream_check() -> None: + if not _parse_bool_env("CORRIDORKEY_AUTO_CHECK_UPSTREAM", True): + return + + global _CHECK_THREAD_STARTED + with _CHECK_THREAD_LOCK: + if _CHECK_THREAD_STARTED: + return + _CHECK_THREAD_STARTED = True + + thread = threading.Thread( + target=_run_upstream_check, + name="CorridorKeyUpstreamCheck", + daemon=True, + ) + thread.start() diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..5b437ba --- /dev/null +++ b/nodes.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +try: + from .corridor_key import CorridorKeyProcessor, CorridorKeySettings +except ImportError: + from corridor_key import CorridorKeyProcessor, CorridorKeySettings + + +class CorridorKey: + @classmethod + def INPUT_TYPES(cls) -> dict[str, dict[str, tuple]]: + return { + "required": { + "image": ("IMAGE",), + "mask": ("MASK",), + "gamma_space": (["sRGB", "Linear"],), + "feather_radius": ( + "INT", + { + "default": 4, + "min": 0, + "max": 64, + "step": 1, + }, + ), + "edge_focus": ( + "FLOAT", + { + "default": 1.5, + "min": 0.0, + "max": 4.0, + "step": 0.1, + }, + ), + "threshold": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + }, + ), + "preserve_core": ( + "FLOAT", + { + "default": 0.2, + "min": 0.0, + "max": 0.45, + "step": 0.01, + }, + ), + "despill_strength": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + }, + ), + "refiner_strength": ( + "FLOAT", + { + "default": 1.0, + "min": 0.1, + "max": 4.0, + "step": 0.1, + }, + ), + "auto_despeckle": (["On", "Off"],), + "despeckle_size": ( + "INT", + { + "default": 400, + "min": 1, + "max": 4096, + "step": 1, + }, + ), + "output_mode": (["Processed", "FG", "Matte", "Comp"],), + } + } + + RETURN_TYPES = ("MASK", "IMAGE", "IMAGE", "IMAGE", "IMAGE") + RETURN_NAMES = ("refined_mask", "selected_output", "foreground", "processed", "comp") + FUNCTION = "run" + CATEGORY = "CorridorKey" + + def __init__(self) -> None: + self._processor = CorridorKeyProcessor() + + def run( + self, + image, + mask, + gamma_space: str, + feather_radius: int, + edge_focus: float, + threshold: float, + preserve_core: float, + despill_strength: float, + refiner_strength: float, + auto_despeckle: str, + despeckle_size: int, + output_mode: str, + ): + settings = CorridorKeySettings( + gamma_space=str(gamma_space), + feather_radius=int(feather_radius), + edge_focus=float(edge_focus), + threshold=float(threshold), + preserve_core=float(preserve_core), + despill_strength=float(despill_strength), + refiner_strength=float(refiner_strength), + auto_despeckle=str(auto_despeckle), + despeckle_size=int(despeckle_size), + output_mode=str(output_mode), + ) + return self._processor.refine(image=image, mask=mask, settings=settings) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..be5cf05 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,46 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "comfyui-corridorkey" +version = "0.1.0" +description = "ComfyUI custom node for CorridorKey-style edge-aware coarse mask refinement." +readme = "README.md" +requires-python = ">=3.10" +dependencies = [ + # Keep runtime deps aligned with requirements.txt for ComfyUI installs. + "numpy>=1.24", + "Pillow>=9.5", + "torch>=2.0", +] + +[project.optional-dependencies] +dev = [ + "black>=24.0.0", + "mypy>=1.10.0", + "pytest>=8.0.0", + "ruff>=0.6.0", +] + +[tool.setuptools] +packages = ["corridor_key"] + +[tool.black] +line-length = 100 +target-version = ["py310"] + +[tool.ruff] +line-length = 100 +target-version = "py310" + +[tool.ruff.lint] +select = ["E", "F", "I", "B"] + +[tool.mypy] +python_version = "3.10" +strict = false +warn_unused_configs = true +disallow_untyped_defs = true +disallow_incomplete_defs = true +no_implicit_optional = true diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..b0926c3 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,7 @@ +# Runtime dependencies for ComfyUI-CorridorKey +# ComfyUI usually already ships with a suitable torch build. Keeping it listed +# here matches common custom-node conventions and lets dependency installers +# verify it is present without changing the node code. +torch>=2.0 +numpy>=1.24 +Pillow>=9.5 diff --git a/tests/test_processor.py b/tests/test_processor.py new file mode 100644 index 0000000..be71806 --- /dev/null +++ b/tests/test_processor.py @@ -0,0 +1,206 @@ +from __future__ import annotations + +import pytest +import torch + +from corridor_key import ( + CorridorKeyProcessor, + CorridorKeySettings, + is_verified_check_conclusions, + select_latest_verified_commit, +) + + +def make_image(batch: int = 1, height: int = 6, width: int = 6) -> torch.Tensor: + image = torch.zeros((batch, height, width, 3), dtype=torch.float32) + image[..., : width // 2, :] = 0.15 + image[..., width // 2 :, :] = 0.85 + return image + + +def test_settings_reject_invalid_threshold() -> None: + with pytest.raises(ValueError, match="threshold"): + CorridorKeySettings(threshold=1.2) + + +def test_settings_match_upstream_like_defaults() -> None: + settings = CorridorKeySettings() + + assert settings.gamma_space == "sRGB" + assert settings.despill_strength == 1.0 + assert settings.refiner_strength == 1.0 + assert settings.auto_despeckle == "On" + assert settings.despeckle_size == 400 + assert settings.output_mode == "Processed" + + +def test_refine_broadcasts_single_mask_across_batch() -> None: + processor = CorridorKeyProcessor() + image = make_image(batch=2) + mask = torch.zeros((6, 6), dtype=torch.float32) + mask[:, 3:] = 1.0 + + refined_mask, selected_output, foreground, processed, comp = processor.refine( + image=image, + mask=mask, + settings=CorridorKeySettings(auto_despeckle="Off"), + ) + + assert refined_mask.shape == (2, 6, 6) + assert selected_output.shape == image.shape + assert foreground.shape == image.shape + assert processed.shape == image.shape + assert comp.shape == image.shape + assert torch.allclose(refined_mask[0], refined_mask[1]) + + +def test_refine_preserves_full_foreground() -> None: + processor = CorridorKeyProcessor() + image = make_image() + mask = torch.ones((1, 6, 6), dtype=torch.float32) + + refined_mask, selected_output, foreground, processed, comp = processor.refine( + image=image, + mask=mask, + settings=CorridorKeySettings(auto_despeckle="Off"), + ) + + assert torch.allclose(refined_mask, torch.ones_like(refined_mask)) + assert torch.allclose(foreground, image) + assert torch.allclose(processed, image) + assert torch.allclose(selected_output, processed) + assert torch.allclose(comp, image) + + +def test_refine_preserves_full_background() -> None: + processor = CorridorKeyProcessor() + image = make_image() + mask = torch.zeros((1, 6, 6), dtype=torch.float32) + + refined_mask, selected_output, foreground, processed, comp = processor.refine( + image=image, + mask=mask, + settings=CorridorKeySettings(auto_despeckle="Off"), + ) + + assert torch.allclose(refined_mask, torch.zeros_like(refined_mask)) + assert torch.allclose(foreground, torch.zeros_like(foreground)) + assert torch.allclose(processed, torch.zeros_like(processed)) + assert not torch.allclose(comp, torch.zeros_like(comp)) + assert torch.allclose(selected_output, processed) + + +def test_refine_rejects_wrong_mask_size() -> None: + processor = CorridorKeyProcessor() + image = make_image(height=6, width=6) + mask = torch.zeros((1, 5, 6), dtype=torch.float32) + + with pytest.raises(ValueError, match="height and width"): + processor.refine( + image=image, + mask=mask, + settings=CorridorKeySettings(auto_despeckle="Off"), + ) + + +def test_output_mode_matte_returns_grayscale_preview() -> None: + processor = CorridorKeyProcessor() + image = make_image() + mask = torch.zeros((1, 6, 6), dtype=torch.float32) + mask[:, :, 3:] = 1.0 + + refined_mask, selected_output, _, _, _ = processor.refine( + image=image, + mask=mask, + settings=CorridorKeySettings(output_mode="Matte", auto_despeckle="Off"), + ) + + assert torch.allclose(selected_output[..., 0], refined_mask) + assert torch.allclose(selected_output[..., 1], refined_mask) + assert torch.allclose(selected_output[..., 2], refined_mask) + + +def test_despill_reduces_green_bias() -> None: + processor = CorridorKeyProcessor() + image = torch.zeros((1, 4, 4, 3), dtype=torch.float32) + image[..., 0] = 0.1 + image[..., 1] = 0.9 + image[..., 2] = 0.1 + mask = torch.ones((1, 4, 4), dtype=torch.float32) + + _, _, foreground_without, _, _ = processor.refine( + image=image, + mask=mask, + settings=CorridorKeySettings(despill_strength=0.0, auto_despeckle="Off"), + ) + _, _, foreground_with, _, _ = processor.refine( + image=image, + mask=mask, + settings=CorridorKeySettings(despill_strength=1.0, auto_despeckle="Off"), + ) + + assert torch.mean(foreground_with[..., 1]) < torch.mean(foreground_without[..., 1]) + + +def test_auto_despeckle_can_remove_single_pixel_island() -> None: + processor = CorridorKeyProcessor() + image = make_image() + mask = torch.zeros((1, 6, 6), dtype=torch.float32) + mask[:, 2, 2] = 1.0 + + refined_mask, _, _, _, _ = processor.refine( + image=image, + mask=mask, + settings=CorridorKeySettings( + auto_despeckle="On", + despeckle_size=9, + feather_radius=0, + preserve_core=0.0, + ), + ) + + assert float(refined_mask[0, 2, 2]) == 0.0 + + +def test_verified_check_policy_requires_success_and_no_failure() -> None: + assert is_verified_check_conclusions(("success", "cancelled")) + assert not is_verified_check_conclusions(("success", "failure")) + assert not is_verified_check_conclusions(("cancelled",)) + + +def test_select_latest_verified_commit_picks_first_clean_candidate() -> None: + commit_payloads = [ + { + "sha": "newest-bad", + "commit": { + "author": {"date": "2026-03-02T00:00:00Z"}, + "message": "Newest bad commit", + }, + }, + { + "sha": "next-good", + "commit": { + "author": {"date": "2026-03-01T00:00:00Z"}, + "message": "Next good commit", + }, + }, + ] + check_payloads_by_sha = { + "newest-bad": { + "check_runs": [ + {"conclusion": "success"}, + {"conclusion": "failure"}, + ] + }, + "next-good": { + "check_runs": [ + {"conclusion": "success"}, + {"conclusion": "cancelled"}, + ] + }, + } + + selected = select_latest_verified_commit(commit_payloads, check_payloads_by_sha) + + assert selected is not None + assert selected.sha == "next-good"