This commit is contained in:
Yousef R. Gamaleldin 2026-07-05 20:42:15 +08:00 committed by GitHub
commit fc09958be6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1101,11 +1101,28 @@ def _encode_image(
For PNG, colorspace selection does not modify pixels PNG is delivered
sRGB-encoded and there is no PNG path for wide-gamut HDR in this node.
"""
if img_tensor.ndim == 2: # 2 channel grayscales
img_tensor = img_tensor.unsqueeze(-1)
height, width, num_channels = img_tensor.shape
has_alpha = num_channels == 4
spec = _FORMAT_SPECS[(file_format, bit_depth, has_alpha)]
frame_fmt = spec["frame_fmt"]
stream_fmt = spec["stream_fmt"]
if num_channels == 1:
if spec["dtype"] == np.float32:
frame_fmt = stream_fmt = "grayf32le"
elif spec["dtype"] == np.uint16:
frame_fmt = "gray16le"
# Ensure big-endian for PNG
stream_fmt = "gray16be" if "be" in stream_fmt else "gray16le"
else:
frame_fmt = stream_fmt = "gray"
if spec["dtype"] == np.float32:
# EXR path: preserve full range, no clamp.
if colorspace == "sRGB":
@ -1124,12 +1141,15 @@ def _encode_image(
codec = av.CodecContext.create(file_format, "w")
codec.width = width
codec.height = height
codec.pix_fmt = spec["stream_fmt"]
codec.pix_fmt = stream_fmt
codec.time_base = Fraction(1, 1)
frame = av.VideoFrame.from_ndarray(img_np, format=spec["frame_fmt"])
if spec["frame_fmt"] != spec["stream_fmt"]:
frame = frame.reformat(format=spec["stream_fmt"])
if num_channels == 1:
img_np = img_np.squeeze(-1)
frame = av.VideoFrame.from_ndarray(img_np, format=frame_fmt)
if frame_fmt != stream_fmt:
frame = frame.reformat(format=stream_fmt)
frame.pts = 0
frame.time_base = codec.time_base