Skip to content

perf: faster JPEG streaming (zero-copy VTK capture, 4:2:0 turbo-jpeg) #65

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Thibault-Pelletier merged 2 commits into from
Oct 9, 2026
Merged

perf: faster JPEG streaming (zero-copy VTK capture, 4:2:0 turbo-jpeg) #65

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9a89ffa
perf: read VTK frames as RGBX without a copy
user27182 Oct 8, 2026
7d30f8a
perf: encode turbo-jpeg with 4:2:0 chroma subsampling
user27182 Oct 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions src/trame_rca/encoders/img.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import numpy as np

TO_IMAGE_TYPE = {
"avif": "image/avif",
"jpeg": "image/jpeg",
Expand All @@ -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.

@Thibault-Pelletier Thibault-Pelletier Oct 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added another comment in the vtk_rca.py, but since the RGBX contiguous view is only required by this function, please move the conversion scope inside this function directly and remove this comment

@user27182 user27182 Oct 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just to confirm, you want this function moved to vtk_rca.py, is that correct?

@Thibault-Pelletier Thibault-Pelletier Oct 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No, I meant that this function should be moved to the turbo_jpeg.py file.

If your change provides any performance update to the PIL code path, then it can stay here, but if not, I would rather this stays an implementation detail of the turbo_jpeg encoding to avoid any unwanted side effects in the PIL case.

@user27182 user27182 Oct 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

rgbx_view applies to both encoders though. For simplicity, if you prefer, I can make this PR Pillow-only, and drop all turbo-jpeg changes, and move rgbx_view into pil.py instead. That should make the implementation and tests a lot clearer in scope.

@Thibault-Pelletier Thibault-Pelletier Oct 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My bad, I had missread the commit messages.
I still think the changes should be scoped to pil and turbo_jpeg encodings as any other changes to the VTK RCA img to row may break the current contract of what the method returns.

So this method can stay here, the comment still needs to be updated.

"""
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
15 changes: 13 additions & 2 deletions src/trame_rca/encoders/pil.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -54,7 +54,18 @@ def encode_np_img_to_bytes(

# t0 = time.time()
fake_file = BytesIO()
image = Image.fromarray(image)
rgbx = rgbx_view(image)
Comment thread
Thibault-Pelletier marked this conversation as resolved.
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")
Expand Down
12 changes: 9 additions & 3 deletions src/trame_rca/encoders/turbo_jpeg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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()
Expand Down Expand Up @@ -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")

Expand Down
7 changes: 6 additions & 1 deletion src/trame_rca/rca/vtk_rca.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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):
Expand All @@ -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]

@Thibault-Pelletier Thibault-Pelletier Oct 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would be the wrong place to put this. If rgbx_view needs a contiguous array, then it can be put in the specific code and not leak here.

@user27182 user27182 Oct 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can drop this and simply do:

np_image = np_image[::-1, :, :3]

Is that acceptable?

@Thibault-Pelletier Thibault-Pelletier Oct 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There shouldn't be any changes to what the VTK side returns as this would break the current contract and break any other users accessing the data and expecting the current format.

If you need this change it would be better to move it inside rgbx_view instead. This way it will only afect this specific code path and not others.

On a side note, I will test your branch on real projects before accepting anything, but I'm concerned the changes in this file may be an hallucination and may flip the render window / flip the color space.

The tests in trame-rca are not remotely complete and just checking a grey cone doesn't inform if you have completely flipped what you wanted to display.

@user27182 user27182 Oct 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can move it.

I understand your concerns. This is why I don't blindly ask Claude to do things without my input and without review first. I only pushed two commits for this PR, but that's because my understanding is kitware does not typically use squash+merge, and so the commits should stand on their own, ready to be integrated as-is directly into the main branch.

And so locally my development branch initially had many commits, iterated upon by me after much back-and-forth with Claude, and then at the end I rewrote history and force-pushed just the two commits. There is lots of local benchmarking and visual testing, as well as testing with trame-pyvista downstream, that is simply not visible here. I also have specific code-review skills I run with Claude on the diff.

If you'd prefer, next time I can add commits without a force-push, and then you can squash+merge into main? What is the standard here for this repo?

@jourdain jourdain Oct 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like to have isolated commit (rather than a big squash), that way the release notes point to each step/goal and it is easier to track. But if a PR is small enough and does not tackle 2+ things at once, then it can be squashed.

For me the split you did in your commits were perfect.

@user27182 user27182 Oct 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There shouldn't be any changes to what the VTK side returns as this would break the current contract and break any other users accessing the data and expecting the current format.

If you need this change it would be better to move it inside rgbx_view instead. This way it will only afect this specific code path and not others.

On a side note, I will test your branch on real projects before accepting anything, but I'm concerned the changes in this file may be an hallucination and may flip the render window / flip the color space.

The tests in trame-rca are not remotely complete and just checking a grey cone doesn't inform if you have completely flipped what you wanted to display.

I added image orientation tests as part of #66. That should help give confidence here that this PR doesn't break anything. Please review #66 first.

@Thibault-Pelletier Thibault-Pelletier Oct 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @user27182, I reviewed #65 and added a small request on git history.
FYI, I tried rebasing the current PR on the other one and there will be a small conflict to address + failing tests due to a too strong coupling of the tests with the actual type of the encoded image.

You may want to relax the tests in #65 to avoid having to fix them just after in this PR.

@user27182 user27182 Oct 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The tests in #66 revealed a strides/transpose bug with #65 here, so I think I'll leave those tests as-is.

The conflict may be unavoidable since both PRs touch overlapping lines in img_cols_rows, probably easiest to deal with the conflict after #66 merges?

@user27182 user27182 Oct 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would be the wrong place to put this. If rgbx_view needs a contiguous array, then it can be put in the specific code and not leak here.

rgbx_view doesn't need a contiguous array. It needs the 4-byte pixels to be in memory top row first, because that's how Pillow and libjpeg-turbo read a buffer in place. np.ascontiguousarray is just how the flip gets done, and master already flips here:

# master: copy the flipped rows back into VTK's own buffer
np_image[:] = np_image[::-1, :, :]

# this PR: copy the flipped rows into a new buffer
np_image = np.ascontiguousarray(np_image[::-1])

Both make the rows top-down. This PR does it with one copy instead of two (NumPy buffers the overlapping in-place assignment through a temporary), and no longer writes into VTK's buffer.

There shouldn't be any changes to what the VTK side returns as this would break the current contract and break any other users accessing the data and expecting the current format.

Agreed that the existing contract should not be broken, and this PR does change the layout of the returned array, though not the values. To let the encoders read 4-byte pixels in place, it captures RGBA and returns the RGB channels as a view ([..., :3]). The shape, dtype, orientation and pixel values are the same as on master, but the array is no longer C-contiguous (4 bytes between pixels instead of 3).

That only matters to code that reads the array's raw memory as packed RGB, e.g. np.frombuffer(arr), arr.view(other_dtype), arr.ctypes.data, or C extensions that require contiguous input. For these cases, an explicit call to np.ascontiguousarray(arr) by user code restores the old layout. Nothing in trame-rca relies on this (the scheduler passes the array straight to the encoders), and the tests from #66 check that every encoder accepts non-contiguous layouts.

The layout change is what makes the zero-copy path possible: with packed RGB, rgbx_view always returns None and Pillow falls back to Image.fromarray, which repacks every pixel. If changing img_cols_rows's layout isn't acceptable, two options:

  1. Keep img_cols_rows exactly as on master, and add a separate method on VtkRemoteControlledArea that returns the RGBA frame, which the scheduler would pass to the encoders when available.
  2. Drop the zero-copy commit and keep only the 4:2:0 one. But this defeats the main purpose of the PR, which is to speed up the default Pillow encoder.

return np_image, cols, rows

@property
Expand Down
Loading