summaryrefslogtreecommitdiff
path: root/mask_tests/mask.py
blob: ee801d4a693d27df69df8ea08298b9d21d23f515 (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
#!/usr/bin/env python3
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import pyopencl as cl
from alive_progress import alive_bar
from matplotlib.backend_bases import MouseButton
from PIL import Image

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 = 0

x_min = (-periods * np.pi) + (square_x * np.pi)
x_max = (periods  * np.pi) + (square_x * np.pi)
y_min = (-periods * np.pi) + (square_y * np.pi)
y_max = (periods * np.pi) + (square_y * np.pi)

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.set_axis_off()
fig.add_axes(ax)

cmap = plt.cm.viridis
cmap.set_bad((0,0,0))
cmap.set_over((0,0,0))
cmap.set_under((0,0,0))

mask_path = "mask.png"

bruh = None

def on_click(event):
  global bruh
  split_ratio = 1
  if (event.button is MouseButton.MIDDLE) and (event.inaxes):
      # there's probably a way to set global coordinates;
      # I don't expect this to go anywhere so I don't really care
      on_x = ((event.xdata / img_res_x) * abs(x_max - x_min)) + x_min
      on_y = ((event.ydata / img_res_y) * abs(y_max - y_min)) + y_min

      #I, uh, also don't know the best way to replicate the openCL code automaticly in python
      #so ajust if nessesary
      x_hops = []
      y_hops = []
      for i in range(iterations):
          x_hops.append(((on_x - x_min) / abs(x_max - x_min)) * img_res_x)
          y_hops.append(((on_y - y_min) / abs(y_max - y_min)) * img_res_y)
          next_x = on_x/np.tan(on_y)
          on_y = on_y/np.tan(on_x)
          on_x = next_x
          if on_x**2 + on_y**2 > escape:
              break
      print(y_hops[0])
      print(on_y)
      print("{} hops".format(len(x_hops)))
      if bruh:
          bruh.pop(0).remove()
      bruh = ax.plot(x_hops, y_hops)
      plt.draw()



      print(on_x, on_y)


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))

#open image
mask = np.asarray(Image.open(mask_path).convert("L"), dtype=np.double)
mask.setflags(write=1)
mask /= np.max(mask) # normalize
print(mask.shape)
print(mask.dtype)
#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)




# TODO clean this up
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,
                                       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="log", 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:
  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()
  ax.imshow(image, norm="log", aspect="auto", cmap=cmap)
  fig.savefig(single_frame_save)
  plt.autoscale(False)
  plt.connect('motion_notify_event', on_click)
  plt.show()