restartingHEADmaster

1 files changed, 165 insertions, 0 deletions

diff --git a/collective/gpu_migration.py b/collective/gpu_migration.py
new file mode 100755
index 0000000..9536561
--- /dev/null
+++ b/collective/gpu_migration.py
@@ -0,0 +1,165 @@
+#!/usr/bin/env python3
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.animation as animation
+import matplotlib.style as mplstyle
+import pyopencl as cl
+from alive_progress import alive_bar
+from matplotlib.backend_bases import MouseButton
+from PIL import Image
+
+# HELLO READER ------------------------------------------------------------------------
+# yes, I know this script is horrid, but it's more of a scratchpad to quickly test ideas
+# and not anything I'm ever going to show anyone else.
+
+img_res_x = 1000
+img_res_y = 1000
+total_pixels = img_res_x * img_res_y # so we don't gotta compute it every time
+
+periods = 1
+square_x = 0
+square_y = -2
+
+escape = 10000
+iterations = (255*3)
+c_x = 2 * np.pi
+c_y = 2 * np.pi
+
+animation_progres_save = "ani1_less.mp4"
+single_frame_save = "asdf.png"
+
+opencl_context = cl.create_some_context(interactive=False)
+opencl_queue = cl.CommandQueue(opencl_context)
+
+kernel_src_path = "./kernel.c"
+
+frames = 1
+rendered_frames = []
+
+image = np.empty([img_res_x, img_res_y], np.double)
+
+plt.style.use('dark_background')
+# fuck this shit
+fig = plt.figure(frameon=False)
+fig.set_size_inches(img_res_x/fig.dpi, img_res_y/fig.dpi)
+#fig.set_size_inches(width/height, 1, forward=False)
+
+
+ax = plt.Axes(fig, [0., 0., 1., 1.])
+
+manual_limits = [15.832664460420503, 
+                 16.161782579162786,
+                 0.21263776938088172,
+                 0.2668613889327035]
+
+
+#ax.invert_yaxis()
+ax.set_xlim((-periods * np.pi) + (square_x * np.pi), (periods  * np.pi) + (square_x * np.pi))
+ax.set_ylim((-periods * np.pi) + (square_y * np.pi), (periods * np.pi) + (square_y * np.pi))
+#ax.set_xlim(manual_limits[0], manual_limits[1])
+#ax.set_ylim(manual_limits[2], manual_limits[3])
+
+#ax.set_axis_off()
+fig.add_axes(ax)
+mplstyle.use('fast')
+
+cmap = plt.cm.viridis
+cmap.set_bad((0,0,0))
+cmap.set_over((0,0,0))
+cmap.set_under((0,0,0))
+
+
+
+
+print("compiling openCL kernel...")
+with open(kernel_src_path, 'r') as kernel_src:
+  compiled_kernel = cl.Program(opencl_context, kernel_src.read()).build()
+
+encoding_progress = alive_bar(frames, bar = 'filling', spinner = 'waves')
+
+def display_encoder_progress(current_frame: int, total_frames: int):
+    print("Encoding: frame {}/{}".format(current_frame, frames))
+
+temp_render_hook = False
+mp_image = None
+
+def render(axes):
+  global temp_render_hook
+  global mp_image
+  if not temp_render_hook:
+    temp_render_hook = True
+    x_min = axes.get_xlim()[0]
+    x_max = axes.get_xlim()[1]
+    y_min = axes.get_ylim()[1]
+    y_max = axes.get_ylim()[0]
+    print(axes.get_ylim())
+    compiled_kernel.render_frame_curvature(opencl_queue, image.shape, None, image_buffer, mask_buffer,
+                                 np.double(abs(x_max - x_min) / img_res_x),
+                                 np.double(abs(y_max - y_min) / img_res_y),
+                                 np.double(x_min), np.double(y_min), 
+                                 np.uint32(iterations), np.uint32(escape),
+                                 np.double(1))
+    cl.enqueue_copy(opencl_queue, image, image_buffer).wait()
+    print("kernel")
+    if mp_image == None:
+      mp_image = ax.imshow(image, norm="linear", aspect="auto", cmap=cmap, interpolation='none', extent=(x_min, x_max, y_min, y_max)) # TODO
+    else:
+      mp_image.set(extent=(x_min, x_max, y_min, y_max))
+      #ax.set_aspect('equal')
+      #mp_image.set_data(image)
+      mp_image.set_array(image)
+      
+
+    fig.canvas.draw_idle()
+    fig.canvas.flush_events()
+    temp_render_hook = False
+
+#plt.ion()
+
+#open image
+mask_path = "mask.png"
+mask = np.asarray(Image.open(mask_path).convert("L").resize((img_res_x, img_res_y)), dtype=np.double)
+#mask = np.zeros((img_res_x, img_res_y), dtype=np.double)
+mask.setflags(write=1)
+mask /= np.max(mask) # normalize
+#print(mask.shape)
+#print(mask.dtype)
+#print(mask)
+#ax.imshow(mask, norm="linear")
+#plt.show()
+
+image_buffer = cl.Buffer(opencl_context, cl.mem_flags.WRITE_ONLY, image.nbytes)
+mask_buffer = cl.Buffer(opencl_context, cl.mem_flags.READ_ONLY, mask.nbytes)
+cl.enqueue_copy(opencl_queue, mask_buffer, mask).wait()
+#mask_buffer = cl.array.Array(opencl_context,mask.shape, dtype=np.uint8, data=mask)
+
+#move to render
+print("Rendering {} frames...".format(frames))
+if frames > 1:
+    x_min = ax.get_xlim()[0]
+    x_max = ax.get_xlim()[1]
+    y_min = ax.get_ylim()[1]
+    y_max = ax.get_ylim()[0]
+    with alive_bar(frames, bar = 'filling', spinner = 'waves') as bar_total:
+        for frame_i in range(0, frames):
+          compiled_kernel.render_frame_curvature(opencl_queue, image.shape, None, image_buffer, mask_buffer,
+                                                 np.double(abs(x_max - x_min) / img_res_x),
+                                                 np.double(abs(y_max - y_min) / img_res_y),
+                                                 np.double(x_min), np.double(y_min), 
+                                                 np.uint32(iterations), np.uint32(escape),
+                                                 np.double(frame_i/frames))
+
+
+          cl.enqueue_copy(opencl_queue, image, image_buffer).wait()
+          rendered_frame = ax.imshow(image, norm="linear", aspect="auto", cmap=cmap, animated="True")
+          rendered_frames.append([rendered_frame])
+          bar_total()
+    print("Encoding/Saving...")
+    ani = animation.ArtistAnimation(fig, rendered_frames, interval=30, blit=True)
+    ani.save(animation_progres_save, extra_args=['-preset', 'lossless'], progress_callback=display_encoder_progress, codec="h264_nvenc")
+else:
+  ax.callbacks.connect('ylim_changed', render)
+  ax.callbacks.connect('xlim_changed', render)
+  render(ax)
+  plt.savefig("out_texture.png")
+  plt.show(block=True)