rig_picker/core/picker.py

759 lines
22 KiB
Python

import bpy
import gpu
from gpu_extras.batch import batch_for_shader
import blf
from mathutils import bvhtree, Vector, Color
from mathutils.geometry import (intersect_point_quad_2d, intersect_point_tri_2d,
intersect_tri_tri_2d)
from ..constants import PICKERS
from .addon_utils import get_operator_from_id
from .geometry_utils import bounding_rect
from pathlib import Path
import re
import json
import os
import threading
class Shape:
def __init__(self, picker, points, polygons=None, edges=None, tooltip='', color=None, source_name=''):
self.type = 'display'
self.picker = picker
self.rig = picker.rig
self.source_name = source_name
self.hover = False
self.press = False
self.shader = gpu.shader.from_builtin('UNIFORM_COLOR')
self.shader.bind()
#self.hover_shader = gpu.shader.from_builtin('UNIFORM_COLOR')
#self.hover_shader.uniform_float("color", [1, 1, 1, 0.1])
self.color = color
self.hover_color = self.brighten_color(self.color)
self.tooltip = tooltip
self.points = points
self.polygons = polygons or []
self.edges = edges or []
self.p_batch = batch_for_shader(self.shader, 'TRIS', {"pos": self.points}, indices=self.polygons)
self.e_batch = batch_for_shader(self.shader, 'LINES', {"pos": self.points}, indices=self.edges)
self.rect = bounding_rect(points)
@property
def color(self):
return self._color
def brighten_color(self, color):
brighten_color = Color(color[:3])
brighten_color.v += 0.05
brighten_color.v *= 1.1
brighten_color.v = max(brighten_color.v, 0.15)
return [*brighten_color, color[3]]
@color.setter
def color(self, color=None):
if not color:
color = [0.5, 0.5, 0.5, 1]
elif isinstance(color, (tuple, list)) and len(color) in (3, 4):
if len(color) == 3:
color = [*color, 1]
elif len(color) == 4:
color = list(color)
else:
raise Exception('color must have a len of 3 or 4')
else:
raise Exception(f'color is {type(color)} must be None or (tuple, list)')
#self.shader.uniform_float("color", color)
self._color = color
def draw(self):
#self.shader.bind()
self.shader.uniform_float("color", self.color)
if self.polygons:
self.p_batch.draw(self.shader)
if self.edges:
self.e_batch.draw(self.shader)
def move_event(self, location):
if not intersect_point_quad_2d(location, *self.rect):
self.hover = False
return False
for p in self.polygons:
if intersect_point_tri_2d(location, *[self.points[i] for i in p]):
self.hover = True
return True
self.hover = False
return False
def press_event(self, mode='SET'):
self.press = True
def release_event(self, mode='SET'):
self.press = False
class BoneShape(Shape):
def __init__(self, picker, points, polygons, edges, bone, tooltip='', color=None, source_name=''):
super().__init__(picker, points=points, polygons=polygons, edges=edges,
tooltip=tooltip, color=color, source_name=source_name)
self.type = 'bone'
self.bone = bone
self.active_color = [1, 1, 1, 0.1]
self.select_color = self.brighten_color(self.hover_color)
self.bone_colors = self.get_bone_colors()
@property
def select(self):
if not self.bone:
return False
return self.bone in (bpy.context.selected_pose_bones or []) #self.bone.bone.select
@property
def active(self):
if not self.bone:
return False
return self.bone == bpy.context.active_pose_bone #self.rig.data.bones.active == self.bone.bone
@property
def hide(self):
if not self.bone:
return False
#return self.bone not in (bpy.context.visible_pose_bones or [])
return self.bone.bone.hide or not any(l.is_visible for l in self.bone.bone.collections)
@property
def bone_color(self):
if not self.bone:
return [0, 0, 0, 1]
if self.select and self.active:
return self.bone_colors['active']
elif self.select:
return self.bone_colors['select']
elif self.hide:
return self.bone_colors['hide']
else:
return self.bone_colors['normal']
def get_bone_colors(self):
theme = bpy.context.preferences.themes['Default']
bone_colors = {
'select': [*theme.view_3d.bone_pose, 1],
'normal': [0.05, 0.05, 0.05, 1],
'active': [*theme.view_3d.bone_pose_active, 1],
'hide': [0.85, 0.85, 0.85, 0.2],
}
if not self.bone:
return bone_colors
if self.bone.color.palette == 'CUSTOM':
bone_color = self.bone
elif self.bone.color.palette == 'DEFAULT':
bone_color = self.bone.bone
normal_color = bone_color.color.custom.normal.copy()
normal_color.s *= 0.75
bone_colors['normal'] = [*normal_color, 1]
bone_colors['select'] = [*bone_color.color.custom.select, 1]
bone_colors['active'] = [*bone_color.color.custom.active, 1]
bone_colors['hide'] = [*normal_color, 0.1]
return bone_colors
def draw_hided(self):
gpu.state.blend_set('ALPHA')
gpu.state.line_width_set(1.0)
line_color = (1, 1, 1, 0.1)
color = Color(self.color[:3])
if self.hover:
color.v += 0.05
color.v *= 1.1
line_color = (1, 1, 1, 0.3)
self.shader.uniform_float("color", (*color[:3], 0.5))
self.p_batch.draw(self.shader)
self.shader.uniform_float("color", line_color)
self.e_batch.draw(self.shader)
gpu.state.blend_set('NONE')
def draw_zero_scaled(self):
gpu.state.blend_set('ALPHA')
gpu.state.line_width_set(1.0)
line_color = (*self.bone_colors['normal'][:3], 0.2)
color = Color(self.color[:3])
if self.hover or self.select:
color.v += 0.05
color.v *= 1.1
line_color = (*self.bone_color[:3], 0.66)
self.shader.uniform_float("color", (*color[:3], 0.5))
self.p_batch.draw(self.shader)
self.shader.uniform_float("color", line_color)
self.e_batch.draw(self.shader)
gpu.state.blend_set('NONE')
def draw(self):
gpu.state.blend_set('ALPHA')
gpu.state.line_width_set(1.0)
if self.hide:
return self.draw_hided()
elif self.bone and self.bone.custom_shape_scale_xyz.length < 0.00001:
return self.draw_zero_scaled()
# Draw Fill
color = self.color
if self.select:
color = self.select_color
elif self.hover:
color = self.hover_color
self.shader.uniform_float("color", color)
self.p_batch.draw(self.shader)
# Draw Outline
gpu.state.blend_set('NONE')
if self.select or self.active:
gpu.state.line_width_set(2.0)
self.shader.uniform_float("color", self.bone_color)
self.e_batch.draw(self.shader)
gpu.state.line_width_set(1.0)
def assign_bone_event(self):
#print('assign_bone_event', self)
scn = bpy.context.scene
rig = self.picker.rig
source_object = scn.objects.get(self.source_name)
if not source_object:
print(f'Source object {self.source_name} not found')
return
active_bone = rig.data.bones.active
if not active_bone:
print('You need to have an active bone')
return
#print(active_bone, source_object)
#print(rig)
source_object.name = rig.data.bones.active.name
source_object.rig_picker.name = rig.data.bones.active.name
def release_event(self, mode='SET'):
super().release_event(mode)
if self.hide or not self.bone:
return
select = True
if mode == 'SUBSTRACT':
select = False
self.bone.bone.select = select
if self.hover:
if mode != 'SUBSTRACT':
self.rig.data.bones.active = self.bone.bone
def border_select(self, border, mode='SET'):
if not self.bone:
return
if self.hide:
self.bone.bone.select = False
return
rect_tri1 = self.rect[0], self.rect[1], self.rect[2]
rect_tri2 = self.rect[2], self.rect[3], self.rect[0]
border_tri1 = border[0], border[1], border[2]
border_tri2 = border[2], border[3], border[0]
if (not intersect_tri_tri_2d(*border_tri1, *rect_tri1) and
not intersect_tri_tri_2d(*border_tri1, *rect_tri2) and
not intersect_tri_tri_2d(*border_tri2, *rect_tri1) and
not intersect_tri_tri_2d(*border_tri2, *rect_tri2)):
return
select = True
if mode == 'SUBSTRACT':
select = False
for polygon in self.polygons:
points = [self.points[i] for i in polygon]
if intersect_tri_tri_2d(*border_tri1, *points):
self.bone.bone.select = select
return
if intersect_tri_tri_2d(*border_tri2, *points):
self.bone.bone.select = select
return
'''
for b in border:
if intersect_point_tri_2d(b, *points):
self.bone.bone.select = select
return
for p in points:
if intersect_point_quad_2d(p, *border):
self.bone.bone.select = select
return
'''
class OperatorShape(Shape):
def __init__(self, picker, points, polygons, operator, tooltip='', color=None, source_name=''):
super().__init__(picker, points=points, polygons=polygons, tooltip=tooltip,
color=color, source_name=source_name)
self.type = 'operator'
self.active_color = [1, 1, 1, 0.15]
self.press_color = [0, 0, 0, 0.25]
self.operator = operator
#self.arguments = arguments#{k:eval(v)}
#self.operator = get_operator_from_id(self.operator)
if not tooltip:
self.tooltip = self.operator.replace('bpy.ops.', '').replace("'INVOKE_DEFAULT', ", '')
#self.reg_args = re.compile(r'(\w+)=')
'''
def parse_args(self):
args = self.reg_args.split(self.arguments)[1:]
#print(args, zip(args[::2], args[1::2]))
return {k: eval(v) for k, v in zip(args[::2], args[1::2])}
#return {k:eval(v) for k, v in self.reg_args.split(self.arguments)}
'''
def release_event(self, mode='SET'):
super().release_event(mode)
#args = self.parse_args()
if not self.operator:
return
exec(self.operator)
#f'bpy.ops;{idname}'
#print(self.idname)
#print(self.arguments)
#else:
# self.bone.bone.select = False
def draw(self):
super().draw()
if self.press:
color = self.press_color
elif self.hover:
color = self.hover_color
else:
return
gpu.state.blend_set('ALPHA')
self.shader.uniform_float("color", color)
self.p_batch.draw(self.shader)
gpu.state.blend_set('NONE')
class Picker:
def __init__(self, parent, rig, shapes):
self.parent = parent
self.region = bpy.context.region
self.rig = rig
self.translation = Vector((0, 0))
self.shapes = []
self.box_select = None
self.hover_shape = None
self.shader = gpu.shader.from_builtin('UNIFORM_COLOR')
self.mouse = None
for shape_data in shapes:
if not shape_data['points']:
continue
if shape_data['type'] in ('CANVAS', 'DISPLAY'):
shape = Shape(
self,
points=shape_data['points'],
polygons=shape_data['polygons'],
edges=shape_data['edges'],
color=shape_data['color']
)
elif shape_data['type'] == 'BONE':
bone = rig.pose.bones.get(shape_data['bone'])
#if not bone:
# print(f'Bone {shape_data["bone"]} not exist')
# continue
shape = BoneShape(
self,
source_name=shape_data['source_name'],
points=shape_data['points'],
polygons=shape_data['polygons'],
edges=shape_data['edges'],
bone=bone,
color=shape_data['color']
)
elif shape_data['type'] == 'OPERATOR':
shape = OperatorShape(
self,
source_name=shape_data['source_name'],
points=shape_data['points'],
polygons=shape_data['polygons'],
operator=shape_data['operator'],
color=shape_data['color'],
tooltip=shape_data['tooltip'],
)
self.shapes.append(shape)
self.rect = bounding_rect([p for s in self.shapes for p in s.points])
def assign_bone_event(self):
for shape in self.shapes:
if shape.type == 'bone' and shape.hover:
shape.assign_bone_event()
bpy.ops.rigpicker.save_picker(index=self.index)
def press_event(self, mode='SET'):
for shape in self.shapes:
#print(s)
if shape.hover:
shape.press_event(mode)
else:
shape.press = False
def release_event(self, mode='SET'):
#bpy.ops.pose.select_all(action='DESELECT')
#print('PICKER release event', mode)
#print(f'type={event.type}, value={event.value}, ctrl={event.ctrl}, shift={event.shift}, alt={event.alt}')
for shape in self.shapes:
if shape.hover:
shape.release_event(mode)
shape.press = False
#bpy.context.area.tag_redraw()
'''
picker.tooltip_event(event='SHOW')
region.tag_redraw()
#context.region.tag_redraw()
picker.tooltip_event(event='HIDE')
bpy.app.timers.register(partial(tooltip, context.region), first_interval=1)
'''
'''
def tooltip_event(self, event):
self.tooltip = ''
if event == 'SHOW':
if self.hover_shape.type == 'bone':
self.tooltip = self.hover_shape.bone.name
#bpy.context.region.tag_redraw()
'''
def border_select(self, border, mode):
border = [Vector(b) - Vector(self.translation) for b in border]
for shape in self.shapes:
shape.press = False
shape.hover = False
if shape.type != 'bone':
continue
shape.border_select(border, mode)
def move_event(self, location):
self.mouse = Vector(location) - Vector(self.translation)
self.hover_shape = None
for shape in reversed(self.shapes):
if self.hover_shape:
shape.hover = False
elif shape.move_event(self.mouse):
self.hover_shape = shape
#if point_inside_rect(self.end, rect):
# over = point_over_shape(self.end,points, edges)
#if bpy.app.timers.is_registered(self.tooltip_event):
#try:
# bpy.app.timers.unregister(self.tooltip_event)
#except:
# pass
#bpy.app.timers.register(self.tooltip_event, first_interval=1)
def draw(self):
with gpu.matrix.push_pop():
gpu.matrix.translate(self.translation)
for shape in self.shapes:
shape.draw()
def under_mouse(self, location):
location = Vector(location) - Vector(self.translation)
return intersect_point_quad_2d(location, *self.rect)
@property
def center(self):
relative_center = (self.rect[1] + self.rect[3]) * 0.5
return relative_center + self.translation
@property
def index(self):
return self.parent.pickers.index(self)
class PickerGroup:
def __init__(self, rig, pickers):
#self.view_location = Vector((0, 0))
self.region = bpy.context.region
self.rig = rig
self.pickers = []
self.mouse = Vector((0, 0))
self.location = Vector((0, 0))
self.hover_shape = None
self.tooltip_shape = None
self.tooltip_mouse = Vector((0, 0))
self.tooltip = ''
self.timer = None
for picker in pickers:
self.add_picker(picker)
def add_picker(self, picker):
self.pickers.append(Picker(self, self.rig, picker))
def draw(self):
y = 0
for picker in self.pickers:
height = picker.rect[1][1] - picker.rect[-1][1]
picker.translation = Vector((0, -y-height*0.5))
picker.draw()
y += height + 50
#break #TODO for now only draw first picker
def move_event(self, mouse):
self.mouse = mouse
self.location = self.region.view2d.region_to_view(*mouse)
# Try to detect view pan to remove tooltip
# view_location = self.region.view2d.region_to_view(0, 0)
# if view_location != self.view_location:
# self.view_location = view_location
# self.tooltip = ''
# if self.timer:
# self.timer.cancel()
# return
hover_shape = None
for picker in self.pickers:
if not picker.under_mouse(self.location):
continue
picker.move_event(self.location)
if picker.hover_shape:
hover_shape = picker.hover_shape
self.hover_shape = hover_shape
if self.tooltip_shape is not self.hover_shape:
self.tooltip = ''
if self.timer:
self.timer.cancel()
self.timer = threading.Timer(0.4, self.tooltip_event)
self.timer.start()
def press_event(self, mode='SET'):
#self.clear_tooltip()
for picker in self.pickers:
if picker.under_mouse(self.location):
picker.press_event(mode)
else:
for shape in picker.shapes:
shape.press = False
def release_event(self, mode='SET'):
if mode == 'SET':
for bone in self.rig.data.bones:
bone.select = False
for picker in self.pickers:
if picker.under_mouse(self.location):
picker.release_event(mode)
else:
for shape in picker.shapes:
shape.press = False
def assign_bone_event(self):
for picker in self.pickers:
if picker.under_mouse(self.location):
picker.assign_bone_event()
def border_select(self, border, mode):
border = [self.region.view2d.region_to_view(*b) for b in border]
if mode == 'SET':
for bone in self.rig.data.bones:
bone.select = False
for picker in self.pickers:
picker.border_select(border, mode)
def clear_tooltip(self):
self.tooltip = ''
self.tooltip_shape = None
self.timer.cancel()
self.region.tag_redraw()
def tooltip_event(self):
#print('Tooltip Event', self)
#print(self.hover_shape, self.hover_shape.type)
if self.hover_shape and self.hover_shape.type != 'display':
if self.hover_shape.type == 'bone' and self.hover_shape.bone:
self.tooltip = self.hover_shape.bone.name
else:
self.tooltip = self.hover_shape.tooltip
self.tooltip_shape = self.hover_shape
else:
return self.clear_tooltip()
self.tooltip_mouse = self.mouse
self.timer.cancel()
#print(self.tooltip)
self.region.tag_redraw()
@property
def rect(self):
return bounding_rect([co - Vector(p.translation) for p in self.pickers for co in p.rect])
@property
def center(self):
center = sum(self.rect, Vector((0, 0))) / len(self.rect)
center[1] = -center[1]
return center
def load_picker_data(rig):
if 'pickers' in rig.data.rig_picker:
picker_datas = [[s.to_dict() for s in p] for p in rig.data.rig_picker['pickers']]
else:
picker_datas = []
for picker in rig.data.rig_picker.sources:
picker_path = Path(bpy.path.abspath(picker.source, library=rig.data.library))
if not picker_path.exists():
print(f'Picker path not exists: {picker_path.resolve()}')
continue
print('Load picker from', picker_path.resolve())
picker_data = json.loads(picker_path.read_text(encoding='utf-8'))
picker_datas.append(picker_data)
PICKERS[rig] = PickerGroup(rig, picker_datas)
def get_picker_path(rig, source, start=None):
picker_path = bpy.path.abspath(source, library=rig.data.library, start=start)
return Path(os.path.abspath(picker_path))
def pack_picker(rig, start=None):
if not 'rig_picker' in rig.data:
return
pickers = []
for picker_source in rig.data['rig_picker'].get('sources', []):
picker_path = get_picker_path(rig, picker_source['source'], start)
if not picker_path.exists():
print(f'{picker_path} not exists')
continue
picker_data = json.loads(picker_path.read_text(encoding='utf-8'))
pickers.append(picker_data)
rig.data['rig_picker']['pickers'] = pickers
def unpack_picker(rig):
if 'rig_picker' not in rig.data.keys():
return
if 'pickers' in rig.data['rig_picker'].keys():
del rig.data['rig_picker']['pickers']