magic_uv/op/uv_sculpt.py

   1 # SPDX-License-Identifier: GPL-2.0-or-later
   2
   3 # <pep8-80 compliant>
   4
   5 __author__ = "Nutti <nutti.metro@gmail.com>"
   6 __status__ = "production"
   7 __version__ = "6.5"
   8 __date__ = "6 Mar 2021"
   9
  10 from math import pi, cos, tan, sin
  11
  12 import bpy
  13 import bmesh
  14 from mathutils import Vector
  15 from bpy_extras import view3d_utils
  16 from mathutils.bvhtree import BVHTree
  17 from mathutils.geometry import barycentric_transform
  18 from bpy.props import (
  19     BoolProperty,
  20     IntProperty,
  21     EnumProperty,
  22     FloatProperty,
  23 )
  24
  25 from .. import common
  26 from ..utils.bl_class_registry import BlClassRegistry
  27 from ..utils.property_class_registry import PropertyClassRegistry
  28 from ..utils import compatibility as compat
  29
  30
  31 if compat.check_version(2, 80, 0) >= 0:
  32     from ..lib import bglx as bgl
  33 else:
  34     import bgl
  35
  36
  37 def _is_valid_context(context):
  38     objs = common.get_uv_editable_objects(context)
  39     if not objs:
  40         return False
  41
  42     # only edit mode is allowed to execute
  43     if context.object.mode != 'EDIT':
  44         return False
  45
  46     # only 'VIEW_3D' space is allowed to execute
  47     if not common.is_valid_space(context, ['VIEW_3D']):
  48         return False
  49
  50     return True
  51
  52
  53 def _get_strength(p, len_, factor):
  54     f = factor
  55
  56     if p > len_:
  57         return 0.0
  58
  59     if p < 0.0:
  60         return f
  61
  62     return (len_ - p) * f / len_
  63
  64
  65 @PropertyClassRegistry()
  66 class _Properties:
  67     idname = "uv_sculpt"
  68
  69     @classmethod
  70     def init_props(cls, scene):
  71         def get_func(_):
  72             return MUV_OT_UVSculpt.is_running(bpy.context)
  73
  74         def set_func(_, __):
  75             pass
  76
  77         def update_func(_, __):
  78             bpy.ops.uv.muv_uv_sculpt('INVOKE_REGION_WIN')
  79
  80         scene.muv_uv_sculpt_enabled = BoolProperty(
  81             name="UV Sculpt",
  82             description="UV Sculpt is enabled",
  83             default=False
  84         )
  85         scene.muv_uv_sculpt_enable = BoolProperty(
  86             name="UV Sculpt Showed",
  87             description="UV Sculpt is enabled",
  88             default=False,
  89             get=get_func,
  90             set=set_func,
  91             update=update_func
  92         )
  93         scene.muv_uv_sculpt_radius = IntProperty(
  94             name="Radius",
  95             description="Radius of the brush",
  96             min=1,
  97             max=500,
  98             default=30
  99         )
 100         scene.muv_uv_sculpt_strength = FloatProperty(
 101             name="Strength",
 102             description="How powerful the effect of the brush when applied",
 103             min=0.0,
 104             max=1.0,
 105             default=0.03,
 106         )
 107         scene.muv_uv_sculpt_tools = EnumProperty(
 108             name="Tools",
 109             description="Select Tools for the UV sculpt brushes",
 110             items=[
 111                 ('GRAB', "Grab", "Grab UVs"),
 112                 ('RELAX', "Relax", "Relax UVs"),
 113                 ('PINCH', "Pinch", "Pinch UVs")
 114             ],
 115             default='GRAB'
 116         )
 117         scene.muv_uv_sculpt_show_brush = BoolProperty(
 118             name="Show Brush",
 119             description="Show Brush",
 120             default=True
 121         )
 122         scene.muv_uv_sculpt_pinch_invert = BoolProperty(
 123             name="Invert",
 124             description="Pinch UV to invert direction",
 125             default=False
 126         )
 127         scene.muv_uv_sculpt_relax_method = EnumProperty(
 128             name="Method",
 129             description="Algorithm used for relaxation",
 130             items=[
 131                 ('HC', "HC", "Use HC method for relaxation"),
 132                 ('LAPLACIAN', "Laplacian",
 133                  "Use laplacian method for relaxation")
 134             ],
 135             default='HC'
 136         )
 137
 138     @classmethod
 139     def del_props(cls, scene):
 140         del scene.muv_uv_sculpt_enabled
 141         del scene.muv_uv_sculpt_enable
 142         del scene.muv_uv_sculpt_radius
 143         del scene.muv_uv_sculpt_strength
 144         del scene.muv_uv_sculpt_tools
 145         del scene.muv_uv_sculpt_show_brush
 146         del scene.muv_uv_sculpt_pinch_invert
 147         del scene.muv_uv_sculpt_relax_method
 148
 149
 150 def location_3d_to_region_2d_extra(region, rv3d, coord):
 151     coord_2d = view3d_utils.location_3d_to_region_2d(region, rv3d, coord)
 152     if coord_2d is None:
 153         prj = rv3d.perspective_matrix @ Vector(
 154             (coord[0], coord[1], coord[2], 1.0))
 155         width_half = region.width / 2.0
 156         height_half = region.height / 2.0
 157         coord_2d = Vector((
 158             width_half + width_half * (prj.x / prj.w),
 159             height_half + height_half * (prj.y / prj.w)
 160         ))
 161     return coord_2d
 162
 163
 164 @BlClassRegistry()
 165 class MUV_OT_UVSculpt(bpy.types.Operator):
 166     """
 167     Operation class: UV Sculpt in View3D
 168     """
 169
 170     bl_idname = "uv.muv_uv_sculpt"
 171     bl_label = "UV Sculpt"
 172     bl_description = "UV Sculpt in View3D"
 173     bl_options = {'REGISTER'}
 174
 175     __handle = None
 176     __timer = None
 177
 178     @classmethod
 179     def poll(cls, context):
 180         # we can not get area/space/region from console
 181         if common.is_console_mode():
 182             return False
 183         return _is_valid_context(context)
 184
 185     @classmethod
 186     def is_running(cls, _):
 187         return 1 if cls.__handle else 0
 188
 189     @classmethod
 190     def handle_add(cls, obj, context):
 191         if not cls.__handle:
 192             sv = bpy.types.SpaceView3D
 193             cls.__handle = sv.draw_handler_add(cls.draw_brush, (obj, context),
 194                                                "WINDOW", "POST_PIXEL")
 195         if not cls.__timer:
 196             cls.__timer = context.window_manager.event_timer_add(
 197                 0.1, window=context.window)
 198             context.window_manager.modal_handler_add(obj)
 199
 200     @classmethod
 201     def handle_remove(cls, context):
 202         if cls.__handle:
 203             sv = bpy.types.SpaceView3D
 204             sv.draw_handler_remove(cls.__handle, "WINDOW")
 205             cls.__handle = None
 206         if cls.__timer:
 207             context.window_manager.event_timer_remove(cls.__timer)
 208             cls.__timer = None
 209
 210     @classmethod
 211     def draw_brush(cls, obj, context):
 212         sc = context.scene
 213         user_prefs = compat.get_user_preferences(context)
 214         prefs = user_prefs.addons["magic_uv"].preferences
 215
 216         num_segment = 180
 217         theta = 2 * pi / num_segment
 218         fact_t = tan(theta)
 219         fact_r = cos(theta)
 220         color = prefs.uv_sculpt_brush_color
 221
 222         bgl.glBegin(bgl.GL_LINE_STRIP)
 223         bgl.glColor4f(color[0], color[1], color[2], color[3])
 224         x = sc.muv_uv_sculpt_radius * cos(0.0)
 225         y = sc.muv_uv_sculpt_radius * sin(0.0)
 226         for _ in range(num_segment):
 227             bgl.glVertex2f(x + obj.current_mco.x, y + obj.current_mco.y)
 228             tx = -y
 229             ty = x
 230             x = x + tx * fact_t
 231             y = y + ty * fact_t
 232             x = x * fact_r
 233             y = y * fact_r
 234         bgl.glEnd()
 235
 236     def __init__(self):
 237         self.__loop_info = {}       # { Object: loop_info }
 238         self.__stroking = False
 239         self.current_mco = Vector((0.0, 0.0))
 240         self.__initial_mco = Vector((0.0, 0.0))
 241
 242     def __stroke_init(self, context, _):
 243         sc = context.scene
 244
 245         self.__initial_mco = self.current_mco
 246
 247         objs = common.get_uv_editable_objects(context)
 248
 249         # get influenced UV
 250         self.__loop_info = {}
 251         for obj in objs:
 252             world_mat = obj.matrix_world
 253             bm = bmesh.from_edit_mesh(obj.data)
 254             uv_layer = bm.loops.layers.uv.verify()
 255             _, region, space = common.get_space('VIEW_3D', 'WINDOW', 'VIEW_3D')
 256
 257             self.__loop_info[obj] = []
 258             for f in bm.faces:
 259                 if not f.select:
 260                     continue
 261                 for i, l in enumerate(f.loops):
 262                     loc_2d = location_3d_to_region_2d_extra(
 263                         region, space.region_3d,
 264                         compat.matmul(world_mat, l.vert.co))
 265                     diff = loc_2d - self.__initial_mco
 266                     if diff.length < sc.muv_uv_sculpt_radius:
 267                         info = {
 268                             "face_idx": f.index,
 269                             "loop_idx": i,
 270                             "initial_vco": l.vert.co.copy(),
 271                             "initial_vco_2d": loc_2d,
 272                             "initial_uv": l[uv_layer].uv.copy(),
 273                             "strength": _get_strength(
 274                                 diff.length, sc.muv_uv_sculpt_radius,
 275                                 sc.muv_uv_sculpt_strength)
 276                         }
 277                         self.__loop_info[obj].append(info)
 278
 279     def __stroke_apply(self, context, _):
 280         sc = context.scene
 281         objs = common.get_uv_editable_objects(context)
 282
 283         for obj in objs:
 284             world_mat = obj.matrix_world
 285             bm = bmesh.from_edit_mesh(obj.data)
 286             uv_layer = bm.loops.layers.uv.verify()
 287             mco = self.current_mco
 288
 289             if sc.muv_uv_sculpt_tools == 'GRAB':
 290                 for info in self.__loop_info[obj]:
 291                     diff_uv = (mco - self.__initial_mco) * info["strength"]
 292                     l = bm.faces[info["face_idx"]].loops[info["loop_idx"]]
 293                     l[uv_layer].uv = info["initial_uv"] + diff_uv / 100.0
 294
 295             elif sc.muv_uv_sculpt_tools == 'PINCH':
 296                 _, region, space = common.get_space(
 297                     'VIEW_3D', 'WINDOW', 'VIEW_3D')
 298                 loop_info = []
 299                 for f in bm.faces:
 300                     if not f.select:
 301                         continue
 302                     for i, l in enumerate(f.loops):
 303                         loc_2d = location_3d_to_region_2d_extra(
 304                             region, space.region_3d,
 305                             compat.matmul(world_mat, l.vert.co))
 306                         diff = loc_2d - self.__initial_mco
 307                         if diff.length < sc.muv_uv_sculpt_radius:
 308                             info = {
 309                                 "face_idx": f.index,
 310                                 "loop_idx": i,
 311                                 "initial_vco": l.vert.co.copy(),
 312                                 "initial_vco_2d": loc_2d,
 313                                 "initial_uv": l[uv_layer].uv.copy(),
 314                                 "strength": _get_strength(
 315                                     diff.length, sc.muv_uv_sculpt_radius,
 316                                     sc.muv_uv_sculpt_strength)
 317                             }
 318                             loop_info.append(info)
 319
 320                 # mouse coordinate to UV coordinate
 321                 ray_vec = view3d_utils.region_2d_to_vector_3d(
 322                     region, space.region_3d, mco)
 323                 ray_vec.normalize()
 324                 ray_orig = view3d_utils.region_2d_to_origin_3d(
 325                     region, space.region_3d, mco)
 326                 ray_tgt = ray_orig + ray_vec * 1000000.0
 327                 mwi = world_mat.inverted()
 328                 ray_orig_obj = compat.matmul(mwi, ray_orig)
 329                 ray_tgt_obj = compat.matmul(mwi, ray_tgt)
 330                 ray_dir_obj = ray_tgt_obj - ray_orig_obj
 331                 ray_dir_obj.normalize()
 332                 tree = BVHTree.FromBMesh(bm)
 333                 loc, _, fidx, _ = tree.ray_cast(ray_orig_obj, ray_dir_obj)
 334                 if not loc:
 335                     return
 336                 loops = [l for l in bm.faces[fidx].loops]
 337                 uvs = [Vector((l[uv_layer].uv.x, l[uv_layer].uv.y, 0.0))
 338                        for l in loops]
 339                 target_uv = barycentric_transform(
 340                     loc,
 341                     loops[0].vert.co, loops[1].vert.co, loops[2].vert.co,
 342                     uvs[0], uvs[1], uvs[2])
 343                 target_uv = Vector((target_uv.x, target_uv.y))
 344
 345                 # move to target UV coordinate
 346                 for info in loop_info:
 347                     l = bm.faces[info["face_idx"]].loops[info["loop_idx"]]
 348                     if sc.muv_uv_sculpt_pinch_invert:
 349                         diff_uv = \
 350                             (l[uv_layer].uv - target_uv) * info["strength"]
 351                     else:
 352                         diff_uv = \
 353                             (target_uv - l[uv_layer].uv) * info["strength"]
 354                     l[uv_layer].uv = l[uv_layer].uv + diff_uv / 10.0
 355
 356             elif sc.muv_uv_sculpt_tools == 'RELAX':
 357                 _, region, space = common.get_space(
 358                     'VIEW_3D', 'WINDOW', 'VIEW_3D')
 359
 360                 # get vertex and loop relation
 361                 vert_db = {}
 362                 for f in bm.faces:
 363                     for l in f.loops:
 364                         if l.vert in vert_db:
 365                             vert_db[l.vert]["loops"].append(l)
 366                         else:
 367                             vert_db[l.vert] = {"loops": [l]}
 368
 369                 # get relaxation information
 370                 for k in vert_db.keys():
 371                     d = vert_db[k]
 372                     d["uv_sum"] = Vector((0.0, 0.0))
 373                     d["uv_count"] = 0
 374
 375                     for l in d["loops"]:
 376                         ln = l.link_loop_next
 377                         lp = l.link_loop_prev
 378                         d["uv_sum"] = d["uv_sum"] + ln[uv_layer].uv
 379                         d["uv_sum"] = d["uv_sum"] + lp[uv_layer].uv
 380                         d["uv_count"] = d["uv_count"] + 2
 381                     d["uv_p"] = d["uv_sum"] / d["uv_count"]
 382                     d["uv_b"] = d["uv_p"] - d["loops"][0][uv_layer].uv
 383                 for k in vert_db.keys():
 384                     d = vert_db[k]
 385                     d["uv_sum_b"] = Vector((0.0, 0.0))
 386                     for l in d["loops"]:
 387                         ln = l.link_loop_next
 388                         lp = l.link_loop_prev
 389                         dn = vert_db[ln.vert]
 390                         dp = vert_db[lp.vert]
 391                         d["uv_sum_b"] = d["uv_sum_b"] + dn["uv_b"] + dp["uv_b"]
 392
 393                 # apply
 394                 for f in bm.faces:
 395                     if not f.select:
 396                         continue
 397                     for i, l in enumerate(f.loops):
 398                         loc_2d = location_3d_to_region_2d_extra(
 399                             region, space.region_3d,
 400                             compat.matmul(world_mat, l.vert.co))
 401                         diff = loc_2d - self.__initial_mco
 402                         if diff.length >= sc.muv_uv_sculpt_radius:
 403                             continue
 404                         db = vert_db[l.vert]
 405                         strength = _get_strength(diff.length,
 406                                                  sc.muv_uv_sculpt_radius,
 407                                                  sc.muv_uv_sculpt_strength)
 408
 409                         base = (1.0 - strength) * l[uv_layer].uv
 410                         if sc.muv_uv_sculpt_relax_method == 'HC':
 411                             t = 0.5 * \
 412                                 (db["uv_b"] + db["uv_sum_b"] / d["uv_count"])
 413                             diff = strength * (db["uv_p"] - t)
 414                             target_uv = base + diff
 415                         elif sc.muv_uv_sculpt_relax_method == 'LAPLACIAN':
 416                             diff = strength * db["uv_p"]
 417                             target_uv = base + diff
 418                         else:
 419                             continue
 420
 421                         l[uv_layer].uv = target_uv
 422
 423             bmesh.update_edit_mesh(obj.data)
 424
 425     def __stroke_exit(self, context, _):
 426         sc = context.scene
 427         objs = common.get_uv_editable_objects(context)
 428
 429         for obj in objs:
 430             bm = bmesh.from_edit_mesh(obj.data)
 431             uv_layer = bm.loops.layers.uv.verify()
 432             mco = self.current_mco
 433
 434             if sc.muv_uv_sculpt_tools == 'GRAB':
 435                 for info in self.__loop_info[obj]:
 436                     diff_uv = (mco - self.__initial_mco) * info["strength"]
 437                     l = bm.faces[info["face_idx"]].loops[info["loop_idx"]]
 438                     l[uv_layer].uv = info["initial_uv"] + diff_uv / 100.0
 439
 440             bmesh.update_edit_mesh(obj.data)
 441
 442     def modal(self, context, event):
 443         if context.area:
 444             context.area.tag_redraw()
 445
 446         if not MUV_OT_UVSculpt.is_running(context):
 447             MUV_OT_UVSculpt.handle_remove(context)
 448             return {'FINISHED'}
 449
 450         self.current_mco = Vector((event.mouse_region_x, event.mouse_region_y))
 451
 452         region_types = [
 453             'HEADER',
 454             'UI',
 455             'TOOLS',
 456             'TOOL_PROPS',
 457         ]
 458         if not common.mouse_on_area(event, 'VIEW_3D') or \
 459            common.mouse_on_regions(event, 'VIEW_3D', region_types):
 460             return {'PASS_THROUGH'}
 461
 462         if event.type == 'LEFTMOUSE':
 463             if event.value == 'PRESS':
 464                 if not self.__stroking:
 465                     self.__stroke_init(context, event)
 466                 self.__stroking = True
 467             elif event.value == 'RELEASE':
 468                 if self.__stroking:
 469                     self.__stroke_exit(context, event)
 470                 self.__stroking = False
 471             return {'RUNNING_MODAL'}
 472         elif event.type == 'MOUSEMOVE':
 473             if self.__stroking:
 474                 self.__stroke_apply(context, event)
 475             return {'RUNNING_MODAL'}
 476         elif event.type == 'TIMER':
 477             if self.__stroking:
 478                 self.__stroke_apply(context, event)
 479             return {'RUNNING_MODAL'}
 480
 481         return {'PASS_THROUGH'}
 482
 483     def invoke(self, context, _):
 484         if context.area:
 485             context.area.tag_redraw()
 486
 487         if MUV_OT_UVSculpt.is_running(context):
 488             MUV_OT_UVSculpt.handle_remove(context)
 489         else:
 490             MUV_OT_UVSculpt.handle_add(self, context)
 491
 492         return {'RUNNING_MODAL'}