diff --git a/nodes/greenscreen.py b/nodes/greenscreen.py index 4392b43..c488f0b 100644 --- a/nodes/greenscreen.py +++ b/nodes/greenscreen.py @@ -10,6 +10,14 @@ from comfy.utils import ProgressBar logger = logging.getLogger("CorridorKey") +def _linear_to_srgb(x: np.ndarray) -> np.ndarray: + return np.where( + x <= 0.0031308, + x * 12.92, + 1.055 * np.power(np.maximum(x, 0.0031308), 1.0 / 2.4) - 0.055, + ) + + class CorridorKeyGreenscreen: @classmethod def INPUT_TYPES(cls): @@ -49,7 +57,7 @@ class CorridorKeyGreenscreen: CATEGORY = "CorridorKey" DESCRIPTION = ( "Neural network greenscreen keying. Produces a clean alpha matte, " - "despilled straight foreground, premultiplied linear RGBA, and a " + "despilled sRGB foreground, straight sRGB RGBA, and a " "checkerboard composite preview." ) @@ -114,11 +122,21 @@ class CorridorKeyGreenscreen: alpha_out = alpha_out[:, :, 0] alphas.append(torch.from_numpy(alpha_out)) - # fg: [H,W,3] sRGB straight float -> IMAGE [H,W,C] - fgs.append(torch.from_numpy(result["fg"])) + # Derive clean despilled foreground from processed RGBA: + # processed is [H,W,4] linear premultiplied — un-premultiply, then + # convert to sRGB to get a foreground that looks like comp_preview + # but without any background baked in. + processed_raw = result["processed"] # [H,W,4] linear premul + proc_alpha = processed_raw[:, :, 3:4] + proc_rgb = processed_raw[:, :, :3] + safe_alpha = np.maximum(proc_alpha, 1e-7) + straight_lin = proc_rgb / safe_alpha + straight_srgb = _linear_to_srgb(np.clip(straight_lin, 0.0, 1.0)) + fgs.append(torch.from_numpy(np.clip(straight_srgb, 0.0, 1.0).astype(np.float32))) - # processed: [H,W,4] linear premultiplied RGBA -> IMAGE [H,W,C] - processed_list.append(torch.from_numpy(result["processed"])) + # processed RGBA: convert to sRGB straight for usable output + proc_srgb = np.concatenate([straight_srgb, proc_alpha], axis=2) + processed_list.append(torch.from_numpy(np.clip(proc_srgb, 0.0, 1.0).astype(np.float32))) # comp: [H,W,3] sRGB checkerboard composite -> IMAGE [H,W,C] comps.append(torch.from_numpy(result["comp"]))