gp_toolbox/OP_eraser_brush.py

import bpy
from bpy.types import Operator
from gpu_extras.presets import draw_circle_2d
from gpu_extras.batch import batch_for_shader
import gpu
from time import time
from mathutils import Vector, Matrix, Euler
from mathutils.kdtree import KDTree
from mathutils.geometry import intersect_line_plane, intersect_line_sphere_2d, intersect_line_line
from bpy_extras.view3d_utils import region_2d_to_location_3d, region_2d_to_vector_3d, \
location_3d_to_region_2d, region_2d_to_origin_3d, region_2d_to_location_3d
from time import time
from math import pi, cos, sin
from .utils import is_locked, is_hidden


def get_gp_mat(gp, name, set_active=False):
    mat = bpy.data.materials.get(name)
    if not mat:
        mat = bpy.data.materials.new(name)
        bpy.data.materials.create_gpencil_data(mat)
        
    if mat not in gp.data.materials[:]:
        gp.data.materials.append(mat)
    mat_index = gp.data.materials[:].index(mat)
    
    if set_active:
        gp.active_material_index = mat_index

    return mat

def get_gp_frame(gp_layer, frame=None):
    if frame is None:
        frame = bpy.context.scene.frame_current
    gp_frame = next((f for f in gp_layer.frames if f.frame_number==frame), None)
    if not gp_frame:
        gp_frame = gp_layer.frames.new(frame)
    
    return gp_frame

def get_gp_layer(gp, name=None):
    if not name:
        return gp.data.layers.active
    
    layer = gp.data.layers.get(name)
    if not layer:
        layer = gp.data.layers.new(name)
        
    gp.data.layers.active = layer

    return layer

def co_2d_to_3d(co, depth=0.1):
    area = bpy.context.area
    region = bpy.context.region
    rv3d = area.spaces.active.region_3d
    view_mat = rv3d.view_matrix.inverted()
    org = view_mat.to_translation()
    
    depth_3d = view_mat @ Vector((0, 0, -depth))
    #org = region_2d_to_origin_3d(region, rv3d, (region.width/2.0, region.height/2.0))
        
    return region_2d_to_location_3d(region, rv3d, co, depth_3d)


    #vec = (region_2d_to_origin_3d(region, rv3d, co) - org).normalized()

    return org + vec

    return org + region_2d_to_vector_3d(region, rv3d , co)

def get_cuts_data(strokes, mouse, radius):
    gp = bpy.context.object
    
    area = bpy.context.area
    region = bpy.context.region
    rv3d = area.spaces.active.region_3d
    view_mat = rv3d.view_matrix.inverted()
    org = view_mat.to_translation()    
    mat = gp.matrix_world
    
    cuts_data = []
    for s in strokes:
        is_polyline = 2<=len(s.points)<=5
        
        if not is_polyline and not s.select:
            continue
        
        print('Cut Stroke', s)
        for i, p in enumerate(s.points):
            
            if not p.select and not is_polyline:
                continue
                # Test if the next or previous is unselected
                
            edges = []
            
            if i > 0:
                prev_p = s.points[i-1]
                if not prev_p.select and not is_polyline:
                    edges.append((i-1, i))
            
            if i < len(s.points)-1:
                next_p = s.points[i+1]
                if not next_p.select or is_polyline:
                    edges.append((i, i+1))
                        
            for p1_index, p2_index in edges:
                p1 = s.points[p1_index]
                p2 = s.points[p2_index]
                
                length_3d = (p2.co-p1.co).length
                
                p1_3d = mat @ p1.co
                p2_3d = mat @ p2.co
                
                p1_2d = location_3d_to_region_2d(region, rv3d, p1_3d)
                p2_2d = location_3d_to_region_2d(region, rv3d, p2_3d)
                
                if p1_2d is None or p2_2d is None:
                    continue
                
                length_2d = (p2_2d-p1_2d).length
                
                if length_2d <= 4:
                    continue
                
                intersects = intersect_line_sphere_2d(p1_2d, p2_2d, mouse, radius+2)
                intersects = [i for i in intersects if i is not None]
                if not intersects:
                    continue
                
                close_points = [(p1_2d-i).length < 1 or (p2_2d-i).length < 1 for i in intersects]
                if any(close_points):
                    #print('close_points', close_points)
                    continue
                    
                
                print('intersects', intersects)

                line_intersects = []
                for i_2d in intersects:                            
                    #factor = ((i_2d-p1_2d).length) / length_2d
                    #factor_3d = factor_2d * length_3d
                    #vec = region_2d_to_vector_3d(region, rv3d, i_2d)
                    #p3_3d = region_2d_to_location_3d(region, rv3d, i_2d, org)
                    #p4_3d = region_2d_to_origin_3d(region, rv3d, i_2d)
                    
                    p3_3d = co_2d_to_3d(i_2d, 0.1)
                    p4_3d = co_2d_to_3d(i_2d, 1000)

                    #bpy.context.scene.cursor.location = p4_3d
                    
                    line_intersect = intersect_line_line(p1_3d, p2_3d, p3_3d, p4_3d)
                    if not line_intersect:
                        continue
                        
                    
                    i1_3d, _ = line_intersect
                    
                    line_intersects += [i1_3d]
                    
                    #context.scene.cursor.location = i1_3d
                
                print('line_intersects', line_intersects)
                if line_intersects:
                    line_intersects.sort(key=lambda x : (x-p1_3d).length)
                    #cut_data[-1].sort(key=lambda x : (x-p1_3d).length)
                    cut_data = [p1_index, p2_index, s, line_intersects]
                    cuts_data.append(cut_data)
        
    return cuts_data


def circle(x, y, radius, segments):
    coords = []
    m = (1.0 / (segments - 1)) * (pi * 2)

    for p in range(segments):
        p1 = x + cos(m * p) * radius
        p2 = y + sin(m * p) * radius
        coords.append((p1, p2))
    return coords


class GPTB_OT_eraser(Operator):
    """Draw a line with the mouse"""
    bl_idname = "gp.eraser"
    bl_label = "Eraser Brush"
    bl_options = {'REGISTER', 'UNDO'}

    def draw_callback_px(self):
        gpu.state.blend_set('ALPHA')
        #bgl.glBlendFunc(bgl.GL_CONSTANT_ALPHA, bgl.GL_ONE_MINUS_CONSTANT_ALPHA)
        #bgl.glBlendColor(1.0, 1.0, 1.0, 0.1)
        
        area = bpy.context.area
        #region = bpy.context.region
        #rv3d = area.spaces.active.region_3d

        bg_color = area.spaces.active.shading.background_color
        #print(bg_color)

        shader = gpu.shader.from_builtin('POLYLINE_UNIFORM_COLOR')
        shader.bind()
        shader.uniform_float("color", (1, 1, 1, 1))
        for mouse, radius in self.mouse_path:
            circle_co = circle(*mouse, radius, 24)
            batch = batch_for_shader(shader, 'TRI_FAN', {"pos":  circle_co})            
            batch.draw(shader)

        draw_circle_2d(self.mouse, (0.75, 0.25, 0.35, 1.0), self.radius, 24)
        gpu.state.blend_set('NONE')

        

    '''
    def draw_holdout(self, context, event):
        gp = context.object
        mat_inv = gp.matrix_world.inverted()
        mouse_3d = co_2d_to_3d(self.mouse)
        radius_3d = co_2d_to_3d(self.mouse + Vector((self.radius, 0)))
        search_radius = (radius_3d-mouse_3d).length
        
        #print('search_radius', search_radius)
        #print('radius', self.radius)
        
        #bpy.context.scene.cursor.location = mouse_3d
        
        for gp_frame, hld_stroke in self.hld_strokes:
            #print('Add Point')
            
            hld_stroke.points.add(count=1)
            p =  hld_stroke.points[-1]
            p.position = mat_inv @ mouse_3d
            p.pressure = search_radius * 2000
            
            #context.scene.cursor.location = mouse_3d
    '''
    def get_radius(self, context, event):
        pressure = event.pressure or 1
        return context.scene.gptoolprops.eraser_radius * pressure
            

        
        
    def erase(self, context, event):
        gp = context.object
        mat_inv = gp.matrix_world.inverted()
            
        new_points = []
        
        #print(self.cuts_data)
        
       # for f in self.gp_frames:
       #     for s in [s for s in f.drawing.strokes if s.material_index==self.hld_index]:
       #         f.drawing.strokes.remove(s)
        
        #gp.data.materials.pop(index=self.hld_index)
        #bpy.data.materials.remove(self.hld_mat)  
        
        
        bpy.ops.object.mode_set(mode='EDIT')
        context.scene.tool_settings.gpencil_selectmode_edit = 'POINT'
        #context.scene.tool_settings.gpencil_selectmode_edit = 'POINT'

        #bpy.ops.gpencil.select_circle(x=x, y=y, radius=radius, wait_for_input=False)
        
        #for cut_data in self.cuts_data:
            
        #    print(cut_data, len(cut_data))
        t0 = time()
        print()
        
        print('Number of cuts', len(self.mouse_path))
        
        for mouse, radius in self.mouse_path:
            t1 = time()
            
            print()
            x, y = mouse
            bpy.ops.gpencil.select_all(action='DESELECT')
            bpy.ops.gpencil.select_circle(x=x, y=y, radius=radius, wait_for_input=False)
            
            strokes = [s for f in self.gp_frames for s in f.drawing.strokes]
            #print('select_circle', time()-t1)
            
            t2 = time()
            cut_data = get_cuts_data(strokes, mouse, radius)
            #print('get_cuts_data', time()-t2)
            
            #print([s for s in strokes if s.select])
            print('cut_data', cut_data)

            t3 = time()
            for p1_index, p2_index, stroke, intersects in cut_data[::-1]:
                bpy.ops.gpencil.select_all(action='DESELECT')
                
                #print('p1_index', p1_index)
                #print('p2_index', p2_index)
                
                p1 = stroke.points[p1_index]
                p2 = stroke.points[p2_index]
                
                p1.select = True
                p2.select = True
                
                number_cuts = len(intersects)            
                
                bpy.ops.gpencil.stroke_subdivide(number_cuts=number_cuts, only_selected=True)
                
                new_p1 = stroke.points[p1_index+1]
                new_p1.position = mat_inv@intersects[0]
                new_points += [(stroke, p1_index+1)]
                
                #print('number_cuts', number_cuts)
                    
                if number_cuts == 2:
                    new_p2 = stroke.points[p1_index+2]
                    new_p2.position = mat_inv@( (intersects[0] + intersects[1])/2 )
                    #new_points += [new_p2]
                
                    new_p3 = stroke.points[p1_index+3]
                    new_p3.position = mat_inv@intersects[1]
                    new_points += [(stroke, p1_index+3)]  
                        
            #print('subdivide', time() - t3)
            
            bpy.ops.gpencil.select_all(action='DESELECT')
            bpy.ops.gpencil.select_circle(x=x, y=y, radius=radius, wait_for_input=False)
            
            '''
            selected_strokes = [s for f in self.gp_frames for s in f.drawing.strokes if s.select]
            tip_points = [p for s in selected_strokes for i, p in enumerate(s.points) if p.select and (i==0 or i == len(s.points)-1)]
            
            bpy.ops.gpencil.select_less()
            
            for p in tip_points:
                p.select = True
            for stroke, index in new_points:
                stroke.points[index].select = False
                '''
            
            t4 = time()
            selected_strokes = [s for f in self.gp_frames for s in f.drawing.strokes if s.select] 

            if selected_strokes:
                bpy.ops.gpencil.delete(type='POINTS')
                
            print('remove points', time()- t4)
            
            
            #print('Total one cut', time()-t1)
            
        #print('Total all cuts', time()-t0)
        #bpy.ops.gpencil.select_less()
        #for stroke, index in new_points:
        #    stroke.points[index].select = False
        
        #bpy.ops.object.mode_set(mode='OBJECT')
        context.scene.tool_settings.gpencil_selectmode_edit = self.gpencil_selectmode_edit
        bpy.ops.object.mode_set(mode='PAINT_GREASE_PENCIL')   
        #selected_strokes = [s for s in self.gp_frame.drawing.strokes if s.select]   
        #bpy.ops.object.mode_set(mode='PAINT_GREASE_PENCIL')
                
    def modal(self, context, event):    
        self.mouse = Vector((event.mouse_region_x, event.mouse_region_y))
        
        self.radius = self.get_radius(context, event)
        context.area.tag_redraw()
        
        if event.type == 'LEFTMOUSE':
            #self.mouse = mouse
            #self.mouse_path.append((self.mouse, self.radius))
            self.erase(context, event)
            bpy.types.SpaceView3D.draw_handler_remove(self._handle, 'WINDOW')
            return {'FINISHED'}
        
        if (self.mouse-self.mouse_prev).length < max(self.radius/1.33, 2):
            return {'RUNNING_MODAL'}    
            
        self.mouse_prev = self.mouse
        
        if event.type in {'MOUSEMOVE', 'INBETWEEN_MOUSEMOVE'}:
            #self.draw_holdout(context, event)
            self.mouse_path.append((self.mouse, self.radius))
            #self.update_cuts_data(context, event)
            #self.erase(context, event)
            return {'RUNNING_MODAL'} 

        elif event.type in {'RIGHTMOUSE', 'ESC'}:
            bpy.types.SpaceView3D.draw_handler_remove(self._handle, 'WINDOW')
            return {'CANCELLED'}
        
        return {'RUNNING_MODAL'}

    def invoke(self, context, event):
        gp = context.object
        matrix = gp.matrix_world
        
        self.gpencil_selectmode_edit = context.scene.tool_settings.gpencil_selectmode_edit
        self.radius = self.get_radius(context, event)
        self.mouse_prev = self.mouse = Vector((event.mouse_region_x, event.mouse_region_y))
        self.mouse_path = [(self.mouse_prev, self.radius)]
        
        area = context.area
        region = context.region
        w, h =  region.width, region.height
        
        rv3d = area.spaces.active.region_3d
        view_mat = rv3d.view_matrix.inverted()
        org = self.org = view_mat.to_translation()
        
        #org = region_2d_to_origin_3d(region, rv3d, (region.width/2.0, region.height/2.0))
        
        
        #print('ORG', org)
        #print('view_mat', view_mat)
        
        
        self.cuts_data = []
        
        #org = self.view_mat @ Vector((0, 0, -10))
        #self.plane_no = self.plane_co-self.org
        
        #bottom_left = region_2d_to_location_3d(region, rv3d , (0, 0), self.plane_co)
        #bottom_right = region_2d_to_location_3d(region, rv3d , (0, w), self.plane_co)
        
        #bottom_left = intersect_line_plane(self.org, bottom_left, self.plane_co, self.plane_no)
        #bottom_right = intersect_line_plane(self.org, bottom_right, self.plane_co, self.plane_no)
        
        #self.scale_fac = (bottom_right-bottom_left).length / w
        
        #print('scale_fac', self.scale_fac)
        #depth_location = view_mat @ Vector((0, 0, -1))
        #context.scene.cursor.location = depth_location
        
        #plane_2d = [(0, 0), (0, h), (w, h), (w, h)]
        #plane_3d = [region_2d_to_location_3d(p)]
        
        
        t0 = time()
        gp_mats = gp.data.materials
        gp_layers = [l for l in gp.data.layers if not is_locked(l) or is_hidden(l)]
        self.gp_frames = [l.current_frame() for l in gp_layers]
        '''
        points_data = [(s, f, gp_mats[s.material_index]) for f in gp_frames for s in f.drawing.strokes]
        points_data = [(s, f, m) for s, f, m in points_data if not m.grease_pencil.hide or m.grease_pencil.lock]
        print('get_gp_points', time()-t0)
        
        t0 = time()
        #points_data = [(s, f, m, p, get_screen_co(p.position, matrix)) for s, f, m in points_data for p in reversed(s.points)]
        points_data = [(s, f, m, p, org + ((matrix @ p.position)-org).normalized()*1) for s, f, m in points_data for p in reversed(s.points)]
        print('points_to_2d', time()-t0)
        
        #print(points_data)
        self.points_data = [(s, f, m, p, co) for s, f, m, p, co in points_data if co is not None]   
        
        #for s, f, m, p, co in self.points_data:
        #    p.position = co
        
        
        t0 = time()
        self.kd_tree = KDTree(len(self.points_data))
        for i, point_data in enumerate(self.points_data):
            s, f, m, p, co = point_data
            self.kd_tree.insert(co, i)
        
        self.kd_tree.balance()
        print('create kdtree', time()-t0)
        '''
        
        '''
        # Create holdout mat
        self.hld_mat = get_gp_mat(gp, name='Eraser Holdout Stroke')
        self.hld_mat.grease_pencil.use_stroke_holdout = True
        self.hld_mat.grease_pencil.show_stroke = True
        self.hld_mat.grease_pencil.show_fill = False
        self.hld_mat.grease_pencil.use_overlap_strokes = True

        self.hld_index = gp_mats[:].index(self.hld_mat)
        
        self.hld_strokes = []
        for f in self.gp_frames:
            hld_stroke = f.drawing.strokes.new()
            hld_stroke.start_cap_mode = 'ROUND'
            hld_stroke.end_cap_mode = 'ROUND'
            hld_stroke.material_index = self.hld_index
            
            #hld_stroke.line_width = self.radius
            
            self.hld_strokes.append((f, hld_stroke))
                          
        self.draw_holdout(context, event)
        '''
        
        
        context.area.tag_redraw()
        self._handle = bpy.types.SpaceView3D.draw_handler_add(self.draw_callback_px, (), 'WINDOW', 'POST_PIXEL')

        context.window_manager.modal_handler_add(self)
                    
        return {'RUNNING_MODAL'}

### --- REGISTER ---

classes=(
GPTB_OT_eraser,
)

def register():
    for cls in classes:
        bpy.utils.register_class(cls)

def unregister():
    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)