hook vtk directly

This commit is contained in:
Felix Pankratz 2025-07-22 19:44:59 +02:00
parent 070c1290b9
commit c074cfce30
2 changed files with 42 additions and 58 deletions

View File

@ -146,7 +146,7 @@ async def main():
plotter.add_mesh(line, color="red", line_width=2) plotter.add_mesh(line, color="red", line_width=2)
#kitty.hide_cursor() kitty.hide_cursor()
try: try:
#asyncio.run(plotter.run()) #asyncio.run(plotter.run())
await plotter.run() await plotter.run()

View File

@ -5,6 +5,11 @@ import select
import sys import sys
import termios import termios
import tty import tty
import vtk
from vtk.util import numpy_support
import numpy as np
from io import BytesIO
from PIL import Image
import pyvista as pv import pyvista as pv
@ -17,13 +22,14 @@ class TerminalPlotter(pv.Plotter):
self.width = width self.width = width
self.height = height self.height = height
self._running = True self._running = True
self._needs_render = True self.w2i_filter = vtk.vtkWindowToImageFilter()
self._render_complete = False self.w2i_filter.SetInputBufferTypeToRGBA()
self._rendering = False self.w2i_filter.ReadFrontBufferOff()
self._render_pending = True self.w2i_filter.SetInput(self.ren_win)
self.set_background([0.0, 0.0, 0.0]) self.set_background([0.0, 0.0, 0.0])
self.add_on_render_callback(self._on_render) # transparency
# if we are too close to the bottom of the terminal, create some space. self.ren_win.SetAlphaBitPlanes(1)
self.ren_win.SetMultiSamples(0)
async def initialize(self): async def initialize(self):
h_pix, _ = await kitty.get_terminal_cell_size() h_pix, _ = await kitty.get_terminal_cell_size()
@ -38,14 +44,15 @@ class TerminalPlotter(pv.Plotter):
print("\n" * self.needed_lines, end="") print("\n" * self.needed_lines, end="")
kitty.set_position(self.start_y, self.start_x) kitty.set_position(self.start_y, self.start_x)
async def _handle_key(self, c):
async def render_to_kitty(self): if c == "a":
self.render() self.camera.Azimuth(10)
buf = io.BytesIO() elif c == "d":
self.screenshot(buf, transparent_background=True, window_size=(self.width, self.height)) self.camera.Azimuth(-10)
await kitty.draw_to_terminal(buf) elif c == "w":
# print("y:", self.start_y, "rows:", self.rows, end="") self.camera.Elevation(10)
kitty.set_position(self.start_y, self.start_x) elif c == "s":
self.camera.Elevation(-10)
async def _input_loop(self): async def _input_loop(self):
fd = sys.stdin.fileno() fd = sys.stdin.fileno()
@ -64,60 +71,37 @@ class TerminalPlotter(pv.Plotter):
finally: finally:
termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
async def _handle_key(self, c):
if c == "a":
self.camera.Azimuth(10)
elif c == "d":
self.camera.Azimuth(-10)
elif c == "w":
self.camera.Elevation(10)
elif c == "s":
self.camera.Elevation(-10)
self._needs_render = True
# await self.render_to_kitty()
def _on_render(self, *args):
"""Callback that runs after each render."""
self._rendering = False
self._render_complete = True
async def _display_loop(self):
"""Loop that captures the screen and updates the terminal after rendering."""
while self._running:
if self._render_complete:
self._render_complete = False
buf = io.BytesIO()
# Safe: Screenshot only after render is confirmed complete
self.screenshot(buf, transparent_background=True, window_size=(self.width, self.height))
await kitty.draw_to_terminal(buf)
kitty.set_position(self.start_y, self.start_x)
await asyncio.sleep(0.001)
async def _render_loop(self): async def _render_loop(self):
while self._running: while self._running:
if self._needs_render and not self._rendering: # keep renedring the scene
self._rendering = True self.ren_win.Render()
self._needs_render = False # Update the filter to grab the current buffer
self.update() # triggers a render and calls _on_render self.w2i_filter.Modified()
await asyncio.sleep(0.001) self.w2i_filter.Update()
async def _display_in_terminal(self, buf): vtk_image = self.w2i_filter.GetOutput()
await kitty.draw_to_terminal(buf)
width, height, _ = vtk_image.GetDimensions()
vtk_array = vtk_image.GetPointData().GetScalars()
arr = numpy_support.vtk_to_numpy(vtk_array)
# Reshape the array to height x width x channels (probably 3 or 4)
arr = arr.reshape(height, width, -1)
# Flip vertically because VTK's origin is bottom-left
arr = np.flip(arr, axis=0)
img = Image.fromarray(arr)
buffer: BytesIO = BytesIO()
img.save(buffer, format="PNG")
await kitty.draw_to_terminal(buffer)
kitty.set_position(self.start_y, self.start_x) kitty.set_position(self.start_y, self.start_x)
async def run(self): async def run(self):
await self.initialize() await self.initialize()
self.iren.initialize() self.iren.initialize()
self._needs_render = True
tasks = [ tasks = [
asyncio.create_task(self._input_loop()), asyncio.create_task(self._input_loop()),
asyncio.create_task(self._render_loop()), asyncio.create_task(self._render_loop()),
asyncio.create_task(self._display_loop()),
] ]
await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
self._running = False self._running = False
#await self.render_to_kitty()
#await self._input_loop()