Fix foreground and processed_rgba outputs to use sRGB straight colorspace

The foreground output was using the engine's raw fg which looked
dark/unusable. Now derived from the processed RGBA by un-premultiplying
and converting linear to sRGB — matches comp_preview quality without
the baked-in checkerboard. The processed_rgba output is also converted
from linear premultiplied to sRGB straight, fixing colorspace and edge
artifacts.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pixelworld AI
2026-05-04 21:16:37 -05:00
co-authored by Claude Opus 4.6
parent 40bc826b01
commit 7eb44e48c1
+23 -5
View File
@@ -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"]))