diff --git a/gpt.py b/gpt.py new file mode 100644 index 0000000..eaa9091 --- /dev/null +++ b/gpt.py @@ -0,0 +1,73 @@ +import pyvista as pv +import vtk +from vtk.util import numpy_support +import numpy as np +import time +import kitty +import asyncio +from PIL import Image +from io import BytesIO + + +async def render_to(image_array): + """ + Dummy rendering function that receives the rendered image as a numpy array. + Replace this with your actual rendering/display logic. + """ + # print(f"Received image with shape: {image_array.shape} and dtype: {image_array.dtype}") + y, x = await kitty.get_position() + img = Image.fromarray(image_array) + buffer: BytesIO = BytesIO() + img.save(buffer, format="PNG") + await kitty.draw_to_terminal(buffer) + kitty.set_position(y, x) + + +async def main(): + plotter = pv.Plotter(off_screen=True, window_size=[400, 400]) + plotter.add_mesh(pv.Sphere(), color="tomato") + plotter.add_text("Live capture example", font_size=14) + + # plotter.show(auto_close=False) # keep window open and interactive + + # Access the vtkRenderWindow + vtk_render_window = plotter.ren_win + + # Setup the window-to-image filter to capture the rendered scene + w2i_filter = vtk.vtkWindowToImageFilter() + w2i_filter.SetInput(vtk_render_window) + + try: + while True: + # Trigger a render to make sure the content is fresh + vtk_render_window.Render() + + # Update the filter to grab the current buffer + w2i_filter.Modified() + w2i_filter.Update() + + vtk_image = w2i_filter.GetOutput() + + 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) + + # Call the dummy render_to function with captured frame + await render_to(arr) + + time.sleep(0.1) # adjust framerate to your needs, e.g. 10 FPS + plotter.camera.Azimuth(1) + + except KeyboardInterrupt: + print("Capture stopped by user.") + plotter.close() + + +if __name__ == "__main__": + asyncio.run(main())