summaryrefslogtreecommitdiff
path: root/polar/gpu_migration.py
blob: 1a8af195848a277694e09e91bc2ed5501c587bcc (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
#!/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)