ee0b6759f13cf27919f506e86753658348b272a7
[blender.git] / doc / python_api / examples / gpu.types.GPUOffScreen.py
1 # Draws an off-screen buffer and display it in the corner of the view.
2 import bpy
3 import bgl
4 import gpu
5 import numpy as np
6
7 g_imageVertSrc = '''
8 in vec2 texCoord;
9 in vec2 pos;
10
11 out vec2 texCoord_interp;
12
13 void main()
14 {
15     gl_Position = vec4(pos.xy, 0.0f, 1.0);
16     gl_Position.z = 1.0f;
17     texCoord_interp = texCoord;
18 }
19 '''
20
21 g_imageFragSrc = '''
22 in vec2 texCoord_interp;
23 out vec4 fragColor;
24
25 uniform sampler2D image;
26
27 void main()
28 {
29     fragColor = texture(image, texCoord_interp);
30 }
31 '''
32
33 g_plane_vertices = np.array([
34     ([-1.0, -1.0], [0.0, 0.0]),
35     ([1.0, -1.0], [1.0, 0.0]),
36     ([1.0, 1.0], [1.0, 1.0]),
37     ([1.0, 1.0], [1.0, 1.0]),
38     ([-1.0, 1.0], [0.0, 1.0]),
39     ([-1.0, -1.0], [0.0, 0.0]),
40 ], [('pos', 'f4', 2), ('uv', 'f4', 2)])
41
42
43 class VIEW3D_OT_draw_offscreen(bpy.types.Operator):
44     bl_idname = "view3d.offscreen_draw"
45     bl_label = "Viewport Offscreen Draw"
46
47     _handle_calc = None
48     _handle_draw = None
49     is_enabled = False
50
51     global_shader = None
52     batch_plane = None
53     uniform_image = -1
54     shader = None
55
56     # manage draw handler
57     @staticmethod
58     def draw_callback_px(self, context):
59         scene = context.scene
60         aspect_ratio = scene.render.resolution_x / scene.render.resolution_y
61
62         self._update_offscreen(context, self._offscreen)
63         self._opengl_draw(context, self._texture, aspect_ratio, 0.2)
64
65     @staticmethod
66     def handle_add(self, context):
67         VIEW3D_OT_draw_offscreen._handle_draw = bpy.types.SpaceView3D.draw_handler_add(
68             self.draw_callback_px, (self, context),
69             'WINDOW', 'POST_PIXEL',
70         )
71
72     @staticmethod
73     def handle_remove():
74         if VIEW3D_OT_draw_offscreen._handle_draw is not None:
75             bpy.types.SpaceView3D.draw_handler_remove(VIEW3D_OT_draw_offscreen._handle_draw, 'WINDOW')
76
77             VIEW3D_OT_draw_offscreen._handle_draw = None
78
79     # off-screen buffer
80     @staticmethod
81     def _setup_offscreen(context):
82         scene = context.scene
83         aspect_ratio = scene.render.resolution_x / scene.render.resolution_y
84
85         try:
86             offscreen = gpu.types.GPUOffScreen(512, int(512 / aspect_ratio))
87         except Exception as ex:
88             print(ex)
89             offscreen = None
90
91         return offscreen
92
93     @staticmethod
94     def _update_offscreen(context, offscreen):
95         scene = context.scene
96         view_layer = context.view_layer
97         render = scene.render
98         camera = scene.camera
99
100         modelview_matrix = camera.matrix_world.inverted()
101         projection_matrix = camera.calc_matrix_camera(
102             context.depsgraph,
103             x=render.resolution_x,
104             y=render.resolution_y,
105             scale_x=render.pixel_aspect_x,
106             scale_y=render.pixel_aspect_y,
107         )
108
109         offscreen.draw_view3d(
110             scene,
111             view_layer,
112             context.space_data,
113             context.region,
114             projection_matrix,
115             modelview_matrix,
116         )
117
118     def _opengl_draw(self, context, texture, aspect_ratio, scale):
119         """
120         OpenGL code to draw a rectangle in the viewport
121         """
122         # view setup
123         bgl.glDisable(bgl.GL_DEPTH_TEST)
124
125         viewport = bgl.Buffer(bgl.GL_INT, 4)
126         bgl.glGetIntegerv(bgl.GL_VIEWPORT, viewport)
127
128         active_texture = bgl.Buffer(bgl.GL_INT, 1)
129         bgl.glGetIntegerv(bgl.GL_TEXTURE_2D, active_texture)
130
131         width = int(scale * viewport[2])
132         height = int(width / aspect_ratio)
133
134         bgl.glViewport(viewport[0], viewport[1], width, height)
135         bgl.glScissor(viewport[0], viewport[1], width, height)
136
137         # draw routine
138         batch_plane = self.get_batch_plane()
139
140         shader = VIEW3D_OT_draw_offscreen.shader
141         # bind it so we can pass the new uniform values
142         shader.bind()
143
144         bgl.glEnable(bgl.GL_TEXTURE_2D)
145         bgl.glActiveTexture(bgl.GL_TEXTURE0)
146         bgl.glBindTexture(bgl.GL_TEXTURE_2D, texture)
147
148         # # TODO, support passing ints
149         # shader.uniform_int(VIEW3D_OT_draw_offscreen.uniform_image, 0)
150         shader.uniform_int("image", 0)
151         batch_plane.draw()
152
153         # restoring settings
154         bgl.glBindTexture(bgl.GL_TEXTURE_2D, active_texture[0])
155         bgl.glDisable(bgl.GL_TEXTURE_2D)
156
157         # reset view
158         bgl.glViewport(viewport[0], viewport[1], viewport[2], viewport[3])
159         bgl.glScissor(viewport[0], viewport[1], viewport[2], viewport[3])
160
161     def get_batch_plane(self):
162         if self.batch_plane is None:
163             global g_plane_vertices
164
165             format = gpu.types.GPUVertFormat()
166             pos_id = format.attr_add(
167                 id="pos",
168                 comp_type='F32',
169                 len=2,
170                 fetch_mode='FLOAT',
171             )
172             uv_id = format.attr_add(
173                 id="texCoord",
174                 comp_type='F32',
175                 len=2,
176                 fetch_mode='FLOAT',
177             )
178             vbo = gpu.types.GPUVertBuf(
179                 len=len(g_plane_vertices),
180                 format=format,
181             )
182
183             vbo.attr_fill(id=pos_id, data=g_plane_vertices["pos"])
184             vbo.attr_fill(id=uv_id, data=g_plane_vertices["uv"])
185
186             batch_plane = gpu.types.GPUBatch(type="TRIS", buf=vbo)
187             shader = self.global_shader
188
189             VIEW3D_OT_draw_offscreen.shader = shader
190             VIEW3D_OT_draw_offscreen.uniform_image = shader.uniform_from_name("image")
191
192             batch_plane.program_set(shader)
193             VIEW3D_OT_draw_offscreen.batch_plane = batch_plane
194         return VIEW3D_OT_draw_offscreen.batch_plane
195
196     # operator functions
197     @classmethod
198     def poll(cls, context):
199         return context.area.type == 'VIEW_3D'
200
201     def modal(self, context, event):
202         if context.area:
203             context.area.tag_redraw()
204
205         if event.type in {'RIGHTMOUSE', 'ESC'}:
206             self.cancel(context)
207             return {'CANCELLED'}
208
209         return {'PASS_THROUGH'}
210
211     def invoke(self, context, event):
212         if VIEW3D_OT_draw_offscreen.is_enabled:
213             self.cancel(context)
214             return {'FINISHED'}
215         else:
216             self._offscreen = VIEW3D_OT_draw_offscreen._setup_offscreen(context)
217             if self._offscreen:
218                 self._texture = self._offscreen.color_texture
219             else:
220                 self.report({'ERROR'}, "Error initializing offscreen buffer. More details in the console")
221                 return {'CANCELLED'}
222
223             VIEW3D_OT_draw_offscreen.handle_add(self, context)
224             VIEW3D_OT_draw_offscreen.is_enabled = True
225
226             if context.area:
227                 context.area.tag_redraw()
228
229             context.window_manager.modal_handler_add(self)
230             return {'RUNNING_MODAL'}
231
232     def cancel(self, context):
233         VIEW3D_OT_draw_offscreen.handle_remove()
234         VIEW3D_OT_draw_offscreen.is_enabled = False
235
236         if VIEW3D_OT_draw_offscreen.batch_plane is not None:
237             del VIEW3D_OT_draw_offscreen.batch_plane
238             VIEW3D_OT_draw_offscreen.batch_plane = None
239
240         VIEW3D_OT_draw_offscreen.shader = None
241
242         if context.area:
243             context.area.tag_redraw()
244
245
246 def register():
247     shader = gpu.types.GPUShader(g_imageVertSrc, g_imageFragSrc)
248     VIEW3D_OT_draw_offscreen.global_shader = shader
249
250     bpy.utils.register_class(VIEW3D_OT_draw_offscreen)
251
252
253 def unregister():
254     bpy.utils.unregister_class(VIEW3D_OT_draw_offscreen)
255     VIEW3D_OT_draw_offscreen.global_shader = None
256
257
258 if __name__ == "__main__":
259     try:
260         unregister()
261     except RuntimeError:
262         pass
263     else:
264         if hasattr(bpy.types, "VIEW3D_OT_draw_offscreen"):
265             del bpy.types.VIEW3D_OT_draw_offscreen.global_shader
266
267     register()