gpu.py

import ctypes
import itertools
import collections
import numpy
import math
import time
import re
from contextlib import contextmanager
import traceback
import contextlib as cl

import bpy
import bgl
from mathutils import Matrix, Vector, geometry

import arnold

from bpy.types import SpaceView3D, SpaceImageEditor, View3DOverlay


class _GPUDRAW():
    def __init__(self):
        self._callback_dict = {}

    def draw(self, width, height, rect):

        v = bgl.Buffer(bgl.GL_FLOAT, 4)
        bgl.glGetFloatv(bgl.GL_VIEWPORT, v)
        vw = v[2]
        vh = v[3]

        import gpu
        import random
        from gpu_extras.batch import batch_for_shader
        from gpu_extras.presets import draw_circle_2d

        vertex_shader = '''
                uniform mat4 modelMatrix;
                uniform mat4 viewProjectionMatrix;

                uniform vec2 fullscreen;
    
                /* Keep in sync with intern/opencolorio/gpu_shader_display_transform_vertex.glsl */
                
                in vec2 texCoord;
                in vec2 pos;
                out vec2 texCoord_interp;

                vec2 normalize_coordinates()
                {
                	return (vec2(2.0) * (pos / fullscreen)) - vec2(1.0);
                }

                void main()
                {
                gl_Position = vec4(normalize_coordinates(), 0.0, 1.0);
                texCoord_interp = texCoord;
                }
        '''
        fragment_shader = '''
            in vec2 texCoord_interp;
            out vec4 fragColor;
            uniform sampler2D image;
            uniform bool ColorMode;
            void main()
            {
                /* if (ColorMode) { */
                fragColor = texture(image, texCoord_interp);
                /* } else { */
                /* fragColor = texture(image, texCoord_interp).rrrr; */
                /* } */
            }
        '''

        def callback_enable(*args, overlay='POST_VIEW'):
            n_id = args[0]
            if n_id in self._callback_dict:
                return

    
            handle_pixel = SpaceView3D.draw_handler_add(draw_callback_px, args, 'WINDOW', overlay)
            SpaceView3D.draw_handler_add(draw_callback_px, args, 'WINDOW', 'PRE_VIEW')
            View3DOverlay.show_edges = True
            View3DOverlay.show_look_dev = True
            self._callback_dict[n_id] = handle_pixel
            #handle_pixel = SpaceView3D.draw_handler_remove(draw_callback_px, args, 'WINDOW', overlay)
            self.tag_redraw_all_3dviews()
        
        def restore_opengl_defaults():
            bgl.glLineWidth(1)
            bgl.glDisable(bgl.GL_BLEND)

        def draw_callback_px(n_id, data):
            context = bpy.context.space_data
            drawing_func = data.get('custom_function')
            args = data.get('args', (None,))       
            drawing_func(args)
            restore_opengl_defaults()
        
        def init_texture(width, height, texname, texture, clr):
            bgl.glEnable(bgl.GL_TEXTURE_2D)
          
            bgl.glActiveTexture(bgl.GL_TEXTURE0)
            
            bgl.glTexParameteri(bgl.GL_TEXTURE_2D, bgl.GL_TEXTURE_MIN_FILTER, bgl.GL_NEAREST)
            bgl.glTexParameteri(bgl.GL_TEXTURE_2D, bgl.GL_TEXTURE_MAG_FILTER, bgl.GL_NEAREST)

            bgl.glTexImage2D(
                bgl.GL_TEXTURE_2D,
                0, bgl.GL_RGBA, width, height,
                0, bgl.GL_RGBA, bgl.GL_FLOAT, texture
            )

        def simple_screen(args):
            """ shader draw function for the texture """

            # border_color = (0.390805, 0.754022, 1.000000, 1.00)
            texture, texname, width, height, batch, shader, cMode = args

            def draw_texture(x=0, y=0, w=width, h=height, texname=texname, c=cMode):
                # function to draw a texture
                bgl.glDisable(bgl.GL_DEPTH_TEST)

                act_tex = bgl.Buffer(bgl.GL_INT, 1)
                bgl.glBindTexture(bgl.GL_TEXTURE_2D, texname)
                
                shader.bind()
                shader.uniform_float("fullscreen", Vector((width, height)))
                shader.uniform_int("image", act_tex)
                #shader.uniform_bool("ColorMode", c)
                batch.draw(shader)

                # restoring settings
                bgl.glBindTexture(bgl.GL_TEXTURE_2D, act_tex[0])
                bgl.glDisable(bgl.GL_TEXTURE_2D)

            draw_texture(x=x, y=y, w=width, h=height, texname=texname, c=cMode)


        def generate_batch_shader(args):
            x, y, w, h = args
            positions = ((x, y), (x + w, y), (x + w, y + h), (x, y + h))
            indices = ((0, 1), (1, 1), (1, 0), (0, 0))
            shader = gpu.types.GPUShader(vertex_shader, fragment_shader)
            batch = batch_for_shader(shader, 'TRI_FAN', {"pos": positions, "texCoord": indices})
            return batch, shader

        # def adjust_position_and_dimensions(x, y, width, height):
        #     """
        #     this could also return scale for a blf notation in the vacinity of the texture
        #     """
        #     scale = 0.5
        #     multiplier = 1.0
        #     #x, y = [x * multiplier, y * multiplier]
        #     #width, height = [width * scale, height * scale]
        #     return x, y, width, height


        total_size = width * height * 3
        texture = bgl.Buffer(bgl.GL_FLOAT, len(rect), rect)
        x = 0
        y = 0
        name = bgl.Buffer(bgl.GL_INT, 1)
        bgl.glGenTextures(1, name)
        #self.texture[n_id] = name[0]
        init_texture(width, height, name[0], texture, 6408)

        #x, y, width, height = adjust_position_and_dimensions(x, y, width, height)
        batch, shader = generate_batch_shader((x, y, width, height))

        is_multi_channel = 1 == 1
        cMode = (is_multi_channel,)

        draw_data = {
            'tree_name': "ARNOLD_WORLD_NODETREE",
            'mode': 'custom_function',
            'custom_function': simple_screen,
            'loc': (0, 0),
            'args': (texture, name[0], width, height, batch, shader, cMode)
        }

        callback_enable(0, draw_data)

    def tag_redraw_all_3dviews(self):
            for window in bpy.context.window_manager.windows:
                for area in window.screen.areas:
                    if area.type == 'VIEW_3D':
                        for region in area.regions:
                            if region.type == 'WINDOW':
                                region.tag_redraw()

    def callback_disable(self, n_id):
        handle_pixel = self._callback_dict.get(n_id, None)
        if not handle_pixel:
            return
        SpaceView3D.draw_handler_remove(handle_pixel, 'WINDOW')
        del self._callback_dict[n_id]
        self.tag_redraw_all_3dviews()