158 lines
5.2 KiB
Python
Executable File
158 lines
5.2 KiB
Python
Executable File
#!/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 = 10000
|
|
img_res_y = 10000
|
|
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.uint32)
|
|
|
|
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.])
|
|
|
|
#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(-6.678564587410841, -4.837591292407222)
|
|
ax.set_ylim(-1.0818908008385455, 0.9287284974589227)
|
|
|
|
#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(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="log", 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:
|
|
with alive_bar(frames, bar = 'filling', spinner = 'waves') as bar_total:
|
|
for frame_i in range(0, frames):
|
|
compiled_kernel.render_frame(opencl_queue, image.shape, None, image_buffer, mask,
|
|
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.png")
|
|
#plt.show(block=True)
|
|
|
|
|