Finalize CorridorKey node, docs, and example workflow
This commit is contained in:
SeanBRVFX
+63
-124
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
@@ -9,21 +10,12 @@ from corridor_key import (
|
||||
is_verified_check_conclusions,
|
||||
select_latest_verified_commit,
|
||||
)
|
||||
from corridor_key.color_utils import composite_straight, linear_to_srgb, srgb_to_linear
|
||||
from corridor_key.tensor_ops import ensure_image_tensor, ensure_mask_batch
|
||||
from nodes import CorridorKey
|
||||
|
||||
|
||||
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:
|
||||
def test_settings_match_upstream_defaults() -> None:
|
||||
settings = CorridorKeySettings()
|
||||
|
||||
assert settings.gamma_space == "sRGB"
|
||||
@@ -31,135 +23,82 @@ def test_settings_match_upstream_like_defaults() -> None:
|
||||
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_settings_reject_invalid_despill() -> None:
|
||||
with pytest.raises(ValueError, match="despill_strength"):
|
||||
CorridorKeySettings(despill_strength=1.5)
|
||||
|
||||
|
||||
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_srgb_linear_round_trip_is_stable() -> None:
|
||||
source = np.array([0.0, 0.018, 0.5, 1.0], dtype=np.float32)
|
||||
round_trip = linear_to_srgb(srgb_to_linear(source))
|
||||
assert np.allclose(round_trip, source, atol=1e-5)
|
||||
|
||||
|
||||
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_composite_straight_matches_expected_formula() -> None:
|
||||
fg = torch.tensor([[[0.8, 0.2, 0.0]]], dtype=torch.float32)
|
||||
bg = torch.tensor([[[0.2, 0.2, 0.2]]], dtype=torch.float32)
|
||||
alpha = torch.tensor([[[0.5]]], dtype=torch.float32)
|
||||
comp = composite_straight(fg, bg, alpha)
|
||||
expected = torch.tensor([[[0.5, 0.2, 0.1]]], dtype=torch.float32)
|
||||
assert torch.allclose(comp, expected)
|
||||
|
||||
|
||||
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()
|
||||
def test_ensure_image_tensor_requires_bhwc() -> None:
|
||||
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)
|
||||
assert ensure_image_tensor(image).shape == image.shape
|
||||
|
||||
_, _, 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])
|
||||
with pytest.raises(ValueError, match="shape"):
|
||||
ensure_image_tensor(torch.zeros((1, 3, 4, 4), dtype=torch.float32))
|
||||
|
||||
|
||||
def test_auto_despeckle_can_remove_single_pixel_island() -> None:
|
||||
def test_ensure_mask_batch_rejects_mismatched_batch() -> None:
|
||||
mask = torch.ones((4, 4), dtype=torch.float32)
|
||||
with pytest.raises(ValueError, match="exactly"):
|
||||
ensure_mask_batch(mask, batch_size=2, height=4, width=4)
|
||||
|
||||
|
||||
def test_node_metadata_populates_ui_info() -> None:
|
||||
input_types = CorridorKey.INPUT_TYPES()
|
||||
required_inputs = input_types["required"]
|
||||
|
||||
assert "coarse" in required_inputs["mask"][1]["tooltip"].lower()
|
||||
assert "per frame" in required_inputs["mask"][1]["tooltip"].lower()
|
||||
assert input_types["hidden"]["unique_id"] == "UNIQUE_ID"
|
||||
assert CorridorKey.RETURN_NAMES == ("fg", "matte", "processed", "QC")
|
||||
assert len(CorridorKey.OUTPUT_TOOLTIPS) == 4
|
||||
assert CorridorKey.DESCRIPTION
|
||||
|
||||
|
||||
def test_processor_reports_progress_messages(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
class FakeEngine:
|
||||
def process_frame(self, **_kwargs):
|
||||
return {
|
||||
"fg": np.zeros((4, 4, 3), dtype=np.float32),
|
||||
"matte": np.zeros((4, 4, 1), dtype=np.float32),
|
||||
"processed": np.zeros((4, 4, 3), dtype=np.float32),
|
||||
"comp": np.zeros((4, 4, 3), dtype=np.float32),
|
||||
}
|
||||
|
||||
monkeypatch.setattr("corridor_key.processor.get_cached_engine", lambda device=None: FakeEngine())
|
||||
|
||||
processor = CorridorKeyProcessor()
|
||||
image = make_image()
|
||||
mask = torch.zeros((1, 6, 6), dtype=torch.float32)
|
||||
mask[:, 2, 2] = 1.0
|
||||
messages: list[tuple[str, int, int]] = []
|
||||
image = torch.zeros((2, 4, 4, 3), dtype=torch.float32)
|
||||
mask = torch.zeros((2, 4, 4), dtype=torch.float32)
|
||||
|
||||
refined_mask, _, _, _, _ = processor.refine(
|
||||
processor.refine(
|
||||
image=image,
|
||||
mask=mask,
|
||||
settings=CorridorKeySettings(
|
||||
auto_despeckle="On",
|
||||
despeckle_size=9,
|
||||
feather_radius=0,
|
||||
preserve_core=0.0,
|
||||
),
|
||||
settings=CorridorKeySettings(),
|
||||
progress_callback=lambda message, completed, total: messages.append((message, completed, total)),
|
||||
)
|
||||
|
||||
assert float(refined_mask[0, 2, 2]) == 0.0
|
||||
assert messages[0] == ("Loading CorridorKey model...", 0, 2)
|
||||
assert messages[1] == ("Processing frame 1/2...", 0, 2)
|
||||
assert messages[2] == ("Processing frame 2/2...", 1, 2)
|
||||
assert messages[-1] == ("CorridorKey complete.", 2, 2)
|
||||
|
||||
|
||||
def test_verified_check_policy_requires_success_and_no_failure() -> None:
|
||||
|
||||
Reference in New Issue
Block a user