kglobe/terminalplotter.py

145 lines
4.9 KiB
Python

import asyncio
import math
import select
import sys
import termios
import tty
import vtk
from vtk.util import numpy_support
import pyvista as pv
import kitty
import numpy as np
class TerminalPlotter(pv.Plotter):
def __init__(self, width, height, **kwargs):
super().__init__(off_screen=True, window_size=(width, height), **kwargs)
self.width = width
self.height = height
self._running = True
# setup vtk rendering chain
self.w2i_filter = vtk.vtkWindowToImageFilter()
self.w2i_filter.SetInputBufferTypeToRGBA()
#self.w2i_filter.ReadFrontBufferOff()
self.w2i_filter.ReadFrontBufferOn()
self.w2i_filter.SetInput(self.ren_win)
# enable transparency
self.set_background([0.0, 0.0, 0.0])
self.ren_win.SetAlphaBitPlanes(1)
self.ren_win.SetMultiSamples(0)
self.terminal = kitty.Terminal()
async def initialize(self):
h_pix, _ = await self.terminal.get_terminal_cell_size()
self.rows, _ = self.terminal.get_terminal_size()
self.start_y, self.start_x = await self.terminal.get_position()
self.needed_lines = math.ceil(self.height / h_pix)
# Add vertical space if needed
if self.rows - self.start_y < self.needed_lines:
missing = self.needed_lines - (self.rows - self.start_y)
self.start_y -= missing
print("\n" * self.needed_lines, end="")
self.terminal.set_position(self.start_y, self.start_x)
self.orbit = self.generate_orbital_path()
def _handle_key(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)
elif c == "+":
self.camera.zoom(1.1)
elif c == "-":
self.camera.zoom(0.9)
elif c == "p":
self.fly_to(self.orbit.points[15])
elif c == "q":
self._running = False
self.terminal.terminal_io.register_input_callback(_handle_key)
# 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)
# elif c == "+":
# self.camera.zoom(1.1)
# elif c == "-":
# self.camera.zoom(0.9)
# elif c == "p":
# self.fly_to(self.orbit.points[15])
#async def _input_loop(self):
# fd = sys.stdin.fileno()
# old_settings = termios.tcgetattr(fd)
# try:
# tty.setcbreak(fd)
# while self._running:
# rlist, _, _ = select.select([sys.stdin], [], [], 0)
# if rlist:
# c = sys.stdin.read(1)
# if c == "q": # quit on q
# self._running = False
# else:
# await self._handle_key(c)
# await asyncio.sleep(0.001)
# finally:
# termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
async def _render_loop(self):
import time
while self._running:
start = time.perf_counter()
# keep renedring the scene
self.ren_win.Render()
# Update the filter to grab the current buffer
self.w2i_filter.Modified()
self.w2i_filter.Update()
vtk_image = self.w2i_filter.GetOutput()
width, height, _ = vtk_image.GetDimensions()
vtk_array = vtk_image.GetPointData().GetScalars()
np_image = numpy_support.vtk_to_numpy(vtk_array).reshape(height, width, 4)
np_image = np_image[::-1] # Flip vertically
np_image = np.ascontiguousarray(np_image) # Ensure memory layout is C-contiguous
#print("Render & copy:", time.perf_counter() - start)
#await self.terminal.draw_to_terminal(np_image, width, height, compress=True)
await self.terminal.draw(np_image, width, height, use_shm=True)
#print("Draw:", time.perf_counter() - start)
self.terminal.set_position(self.start_y, self.start_x)
self.camera.Azimuth(1)
async def run(self):
await self.initialize()
self.iren.initialize()
#tasks = [
# #asyncio.create_task(self._input_loop()),
# asyncio.create_task(self._render_loop()),
#]
#await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
await self._render_loop()
self._running = False
async def main():
plotter = TerminalPlotter(1000, 1000)
mesh = pv.Sphere()
plotter.add_mesh(mesh)
await plotter.run()
if __name__ == '__main__':
import cProfile
cProfile.run("asyncio.run(main())", "profile")