diff --git a/src/trame_rca/encoders/img.py b/src/trame_rca/encoders/img.py index 7f45d74..512c546 100644 --- a/src/trame_rca/encoders/img.py +++ b/src/trame_rca/encoders/img.py @@ -1,3 +1,5 @@ +import numpy as np + TO_IMAGE_TYPE = { "avif": "image/avif", "jpeg": "image/jpeg", @@ -13,3 +15,24 @@ "png": "png", "webp": "webp", } + + +def rgbx_view(image): + """ + Return the RGBX buffer behind an RGB view over 4-byte pixels, or None. + VtkRemoteControlledArea returns such a view so the encoders can skip a copy. + """ + base = image.base + if ( + image.dtype == np.uint8 + and image.ndim == 3 + and image.shape[2] == 3 + and isinstance(base, np.ndarray) + and base.dtype == np.uint8 + and base.flags.c_contiguous + and base.shape == (*image.shape[:2], 4) + and image.strides == base.strides + and image.ctypes.data == base.ctypes.data + ): + return base + return None diff --git a/src/trame_rca/encoders/pil.py b/src/trame_rca/encoders/pil.py index 9f4f99c..d1db731 100644 --- a/src/trame_rca/encoders/pil.py +++ b/src/trame_rca/encoders/pil.py @@ -4,7 +4,7 @@ from numpy.typing import NDArray from PIL import Image -from trame_rca.encoders.img import TO_IMAGE_FORMAT, TO_IMAGE_TYPE +from trame_rca.encoders.img import TO_IMAGE_FORMAT, TO_IMAGE_TYPE, rgbx_view # Not supported on all platform so make it optional with suppress(ImportError): @@ -54,7 +54,18 @@ def encode_np_img_to_bytes( # t0 = time.time() fake_file = BytesIO() - image = Image.fromarray(image) + rgbx = rgbx_view(image) + if rgbx is None: + # Copies: Pillow stores RGB as 4 bytes per pixel, so it can't share a + # 3-channel array + image = Image.fromarray(image) + else: + # VtkRemoteControlledArea frame: map its RGBX buffer without a copy + height, width = rgbx.shape[:2] + image = Image.frombuffer("RGBX", (width, height), rgbx, "raw", "RGBX", 0, 1) + if TO_IMAGE_FORMAT[img_format] != "jpeg": + # PNG can't write RGBX; this is the one copy fromarray would make + image = image.convert("RGB") image.save(fake_file, TO_IMAGE_FORMAT[img_format], quality=quality) # t1 = time.time() # print(f"pill encode {t1-t0:.04f}s") diff --git a/src/trame_rca/encoders/turbo_jpeg.py b/src/trame_rca/encoders/turbo_jpeg.py index 094033c..1b836b0 100644 --- a/src/trame_rca/encoders/turbo_jpeg.py +++ b/src/trame_rca/encoders/turbo_jpeg.py @@ -1,6 +1,6 @@ from numpy.typing import NDArray -from turbojpeg import TurboJPEG, TJPF_RGB -from trame_rca.encoders.img import TO_IMAGE_TYPE +from turbojpeg import TurboJPEG, TJPF_RGB, TJPF_RGBX, TJSAMP_420 +from trame_rca.encoders.img import TO_IMAGE_TYPE, rgbx_view # import time jpeg = TurboJPEG() @@ -41,7 +41,13 @@ def encode_np_img_to_bytes( return b"" # t0 = time.time() - result = jpeg.encode(image, quality=quality, pixel_format=TJPF_RGB) + pixel_format = TJPF_RGB + rgbx = rgbx_view(image) + if rgbx is not None: + image, pixel_format = rgbx, TJPF_RGBX + result = jpeg.encode( + image, quality=quality, pixel_format=pixel_format, jpeg_subsample=TJSAMP_420 + ) # t1 = time.time() # print(f"tubo-jpeg encode {t1-t0:.04f}s") diff --git a/src/trame_rca/rca/vtk_rca.py b/src/trame_rca/rca/vtk_rca.py index 534dcaa..06360e3 100644 --- a/src/trame_rca/rca/vtk_rca.py +++ b/src/trame_rca/rca/vtk_rca.py @@ -1,5 +1,6 @@ import json +import numpy as np import vtkmodules.vtkRenderingOpenGL2 # noqa from packaging.version import Version from trame_common.utils import profiler @@ -27,6 +28,7 @@ def __init__(self, vtk_render_window: vtkRenderWindow): self._window_to_image.ReadFrontBufferOff() self._window_to_image.ShouldRerenderOff() self._window_to_image.FixBoundaryOn() + self._window_to_image.SetInputBufferTypeToRGBA() @property def img_cols_rows(self): @@ -40,7 +42,10 @@ def img_cols_rows(self): scalars = image_data.GetPointData().GetScalars() np_image = vtk_to_numpy(scalars) np_image = np_image.reshape((rows, cols, -1)) - np_image[:] = np_image[::-1, :, :] + # Flip into a new RGBA buffer and return its RGB channels as a view: + # VTK's alpha is 0 on the background, and the encoders can read the + # 4-byte pixels in place (see encoders.img.rgbx_view) + np_image = np.ascontiguousarray(np_image[::-1])[..., :3] return np_image, cols, rows @property