vdm_brush_baker/__init__.py

   1 # SPDX-FileCopyrightText: 2023 Robin Hohnsbeen
   2 #
   3 # SPDX-License-Identifier: GPL-3.0-or-later
   4
   5 from mathutils import Vector
   6 from datetime import datetime
   7 from pathlib import Path
   8 import os
   9 import bpy
  10 from . import bakematerial
  11 import importlib
  12 importlib.reload(bakematerial)
  13
  14 bl_info = {
  15     'name': 'VDM Brush Baker',
  16     'author': 'Robin Hohnsbeen',
  17     'description': 'Bake vector displacement brushes easily from a plane',
  18     'blender': (3, 5, 0),
  19     'version': (1, 0, 2),
  20     'location': 'Sculpt Mode: View3D > Sidebar > Tool Tab',
  21     'warning': '',
  22     'category': 'Baking',
  23     'doc_url': '{BLENDER_MANUAL_URL}/addons/baking/vdm_brush_baker.html'
  24 }
  25
  26
  27 class vdm_brush_baker_addon_data(bpy.types.PropertyGroup):
  28     draft_brush_name: bpy.props.StringProperty(
  29         name='Name',
  30         default='NewVDMBrush',
  31         description='The name that will be used for the brush and texture')
  32     render_resolution: bpy.props.EnumProperty(items={
  33         ('128', '128 px', 'Render with 128 x 128 pixels', 1),
  34         ('256', '256 px', 'Render with 256 x 256 pixels', 2),
  35         ('512', '512 px', 'Render with 512 x 512 pixels', 3),
  36         ('1024', '1024 px', 'Render with 1024 x 1024 pixels', 4),
  37         ('2048', '2048 px', 'Render with 2048 x 2048 pixels', 5),
  38     },
  39         default='512', name='Map Resolution')
  40     compression: bpy.props.EnumProperty(items={
  41         ('none', 'None', '', 1),
  42         ('zip', 'ZIP (lossless)', '', 2),
  43     },
  44         default='zip', name='Compression')
  45     color_depth: bpy.props.EnumProperty(items={
  46         ('16', '16', '', 1),
  47         ('32', '32', '', 2),
  48     },
  49         default='16',
  50         name='Color Depth',
  51         description='A color depth of 32 can give better results but leads to far bigger file sizes. 16 should be good if the sculpt doesn\'t extend "too far" from the original plane')
  52     render_samples: bpy.props.IntProperty(name='Render Samples',
  53                                           default=64,
  54                                           min=2,
  55                                           max=4096)
  56
  57
  58 def get_addon_data() -> vdm_brush_baker_addon_data:
  59     return bpy.context.scene.VDMBrushBakerAddonData
  60
  61
  62 def get_output_path(filename):
  63     save_path = bpy.path.abspath('/tmp')
  64     if bpy.data.is_saved:
  65         save_path = os.path.dirname(bpy.data.filepath)
  66     save_path = os.path.join(save_path, 'output_vdm', filename)
  67
  68     if bpy.data.is_saved:
  69         return bpy.path.relpath(save_path)
  70     else:
  71         return save_path
  72
  73
  74 class PT_VDMBaker(bpy.types.Panel):
  75     """
  76     The main panel of the add-on which contains a button to create a sculpting plane and a button to bake the vector displacement map.
  77     It also has settings for name (image, texture and brush at once), resolution, compression and color depth.
  78     """
  79     bl_label = 'VDM Brush Baker'
  80     bl_idname = 'Editor_PT_LayoutPanel'
  81     bl_space_type = 'VIEW_3D'
  82     bl_region_type = 'UI'
  83     bl_category = 'Tool'
  84
  85     def draw(self, context):
  86         layout = self.layout
  87         addon_data = get_addon_data()
  88
  89         layout.use_property_split = True
  90         layout.use_property_decorate = False
  91
  92         layout.operator(create_sculpt_plane.bl_idname, icon='ADD')
  93
  94         layout.separator()
  95
  96         is_occupied, brush_name = get_new_brush_name()
  97         button_text = 'Overwrite VDM Brush' if is_occupied else 'Render and Create VDM Brush'
  98
  99         createvdmlayout = layout.row()
 100         createvdmlayout.enabled = context.active_object is not None and context.active_object.type == 'MESH'
 101         createvdmlayout.operator(
 102             create_vdm_brush.bl_idname, text=button_text, icon='BRUSH_DATA')
 103
 104         if is_occupied:
 105             layout.label(
 106                 text='Name Taken: Brush will be overwritten.', icon='INFO')
 107
 108         col = layout.column()
 109         col.alert = is_occupied
 110         col.prop(addon_data, 'draft_brush_name')
 111
 112         settings_layout = layout.column(align=True)
 113         settings_layout.label(text='Settings')
 114         layout_box = settings_layout.box()
 115
 116         col = layout_box.column()
 117         col.prop(addon_data, 'render_resolution')
 118
 119         col = layout_box.column()
 120         col.prop(addon_data, 'compression')
 121
 122         col = layout_box.column()
 123         col.prop(addon_data, 'color_depth')
 124
 125         col = layout_box.column()
 126         col.prop(addon_data, 'render_samples')
 127
 128         layout.separator()
 129
 130
 131 def get_new_brush_name():
 132     addon_data = get_addon_data()
 133
 134     is_name_occupied = False
 135     for custom_brush in bpy.data.brushes:
 136         if custom_brush.name == addon_data.draft_brush_name:
 137             is_name_occupied = True
 138             break
 139
 140     if addon_data.draft_brush_name != '':
 141         return is_name_occupied, addon_data.draft_brush_name
 142     else:
 143         date = datetime.now()
 144         dateformat = date.strftime('%b-%d-%Y-%H-%M-%S')
 145         return False, f'vdm-{dateformat}'
 146
 147
 148 class create_sculpt_plane(bpy.types.Operator):
 149     """
 150     Creates a grid with 128 vertices per side plus two multires subdivisions.
 151     It uses 'Preserve corners' so further subdivisions can be made while the corners of the grid stay pointy.
 152     """
 153     bl_idname = 'sculptplane.create'
 154     bl_label = 'Create Sculpting Plane'
 155     bl_description = 'Creates a plane with a multires modifier to sculpt on'
 156     bl_options = {'REGISTER', 'UNDO'}
 157
 158     def execute(self, context):
 159         bpy.ops.mesh.primitive_grid_add(x_subdivisions=128, y_subdivisions=128, size=2,
 160                                         enter_editmode=False, align='WORLD', location=(0, 0, 0), scale=(1, 1, 1))
 161         new_grid = bpy.context.active_object
 162         multires = new_grid.modifiers.new('MultiresVDM', type='MULTIRES')
 163         multires.boundary_smooth = 'PRESERVE_CORNERS'
 164         bpy.ops.object.multires_subdivide(
 165             modifier='MultiresVDM', mode='CATMULL_CLARK')
 166         bpy.ops.object.multires_subdivide(
 167             modifier='MultiresVDM', mode='CATMULL_CLARK')  # 512 vertices per one side
 168
 169         return {'FINISHED'}
 170
 171
 172 class create_vdm_brush(bpy.types.Operator):
 173     """
 174     This operator will bake a vector displacement map from the active object and create a texture and brush datablock.
 175     """
 176     bl_idname = 'vdmbrush.create'
 177     bl_label = 'Create vdm brush from plane'
 178     bl_description = 'Creates a vector displacement map from your sculpture and creates a brush with it'
 179     bl_options = {'REGISTER', 'UNDO'}
 180
 181     @classmethod
 182     def poll(cls, context):
 183         return context.active_object is not None and context.active_object.type == 'MESH'
 184
 185     def execute(self, context):
 186         if context.active_object is None or context.active_object.type != 'MESH':
 187             return {'CANCELLED'}
 188
 189         vdm_plane = context.active_object
 190
 191         addon_data = get_addon_data()
 192         new_brush_name = addon_data.draft_brush_name
 193         reference_brush_name = addon_data.draft_brush_name
 194
 195         is_occupied, brush_name = get_new_brush_name()
 196         if len(addon_data.draft_brush_name) == 0 or is_occupied:
 197             addon_data.draft_brush_name = brush_name
 198
 199         # Saving user settings
 200         scene = context.scene
 201         default_render_engine = scene.render.engine
 202         default_view_transform = scene.view_settings.view_transform
 203         default_display_device = scene.display_settings.display_device
 204         default_file_format = scene.render.image_settings.file_format
 205         default_color_mode = scene.render.image_settings.color_mode
 206         default_codec = scene.render.image_settings.exr_codec
 207         default_denoise = scene.cycles.use_denoising
 208         default_compute_device = scene.cycles.device
 209         default_scene_samples = scene.cycles.samples
 210         default_plane_location = vdm_plane.location.copy()
 211         default_plane_rotation = vdm_plane.rotation_euler.copy()
 212         default_mode = bpy.context.object.mode
 213
 214         bpy.ops.object.mode_set(mode='OBJECT')
 215
 216         vdm_bake_material = bakematerial.get_vdm_bake_material()
 217         try:
 218             # Prepare baking
 219             scene.render.engine = 'CYCLES'
 220             scene.cycles.samples = addon_data.render_samples
 221             scene.cycles.use_denoising = False
 222             scene.cycles.device = 'GPU'
 223
 224             old_image_name = f'{reference_brush_name}'
 225             if old_image_name in bpy.data.images:
 226                 bpy.data.images[old_image_name].name = 'Old VDM texture'
 227
 228                 # Removing the image right away can cause a crash when sculpt mode is exited.
 229                 # bpy.data.images.remove(bpy.data.images[old_image_name])
 230
 231             vdm_plane.data.materials.clear()
 232             vdm_plane.data.materials.append(vdm_bake_material)
 233             vdm_plane.location = Vector([0, 0, 0])
 234             vdm_plane.rotation_euler = (0, 0, 0)
 235
 236             vdm_texture_node = vdm_bake_material.node_tree.nodes['VDMTexture']
 237             render_resolution = int(addon_data.render_resolution)
 238
 239             bpy.ops.object.select_all(action='DESELECT')
 240             vdm_plane.select_set(True)
 241             output_path = get_output_path(f'{new_brush_name}.exr')
 242             vdm_texture_image = bpy.data.images.new(
 243                 name=new_brush_name, width=render_resolution, height=render_resolution, alpha=False, float_buffer=True)
 244             vdm_bake_material.node_tree.nodes.active = vdm_texture_node
 245
 246             vdm_texture_image.filepath_raw = output_path
 247
 248             scene.render.image_settings.file_format = 'OPEN_EXR'
 249             scene.render.image_settings.color_mode = 'RGB'
 250             scene.render.image_settings.exr_codec = 'NONE'
 251             if addon_data.compression == 'zip':
 252                 scene.render.image_settings.exr_codec = 'ZIP'
 253
 254             scene.render.image_settings.color_depth = addon_data.color_depth
 255             vdm_texture_image.use_half_precision = addon_data.color_depth == '16'
 256
 257             vdm_texture_image.colorspace_settings.is_data = True
 258             vdm_texture_image.colorspace_settings.name = 'Non-Color'
 259
 260             vdm_texture_node.image = vdm_texture_image
 261             vdm_texture_node.select = True
 262
 263             # Bake
 264             bpy.ops.object.bake(type='EMIT')
 265             # save as render so we have more control over compression settings
 266             vdm_texture_image.save_render(
 267                 filepath=bpy.path.abspath(output_path), scene=scene, quality=0)
 268             # Removes the dirty flag, so the image doesn't have to be saved again by the user.
 269             vdm_texture_image.pack()
 270             vdm_texture_image.unpack(method='REMOVE')
 271
 272         except BaseException as Err:
 273             self.report({"ERROR"}, f"{Err}")
 274
 275         finally:
 276             scene.render.image_settings.file_format = default_file_format
 277             scene.render.image_settings.color_mode = default_color_mode
 278             scene.render.image_settings.exr_codec = default_codec
 279             scene.cycles.samples = default_scene_samples
 280             scene.display_settings.display_device = default_display_device
 281             scene.view_settings.view_transform = default_view_transform
 282             scene.cycles.use_denoising = default_denoise
 283             scene.cycles.device = default_compute_device
 284             scene.render.engine = default_render_engine
 285             vdm_plane.data.materials.clear()
 286             vdm_plane.location = default_plane_location
 287             vdm_plane.rotation_euler = default_plane_rotation
 288
 289         # Needs to be in sculpt mode to set 'AREA_PLANE' mapping on new brush.
 290         bpy.ops.object.mode_set(mode='SCULPT')
 291
 292         # Texture
 293         vdm_texture: bpy.types.Texture = None
 294         if bpy.data.textures.find(reference_brush_name) != -1:
 295             vdm_texture = bpy.data.textures[reference_brush_name]
 296         else:
 297             vdm_texture = bpy.data.textures.new(
 298                 name=new_brush_name, type='IMAGE')
 299         vdm_texture.extension = 'EXTEND'
 300         vdm_texture.image = vdm_texture_image
 301         vdm_texture.name = new_brush_name
 302
 303         # Brush
 304         new_brush: bpy.types.Brush = None
 305         if bpy.data.brushes.find(reference_brush_name) != -1:
 306             new_brush = bpy.data.brushes[reference_brush_name]
 307             self.report({'INFO'}, f'Changed draw brush \'{new_brush.name}\'')
 308         else:
 309             new_brush = bpy.data.brushes.new(
 310                 name=new_brush_name, mode='SCULPT')
 311             self.report(
 312                 {'INFO'}, f'Created new draw brush \'{new_brush.name}\'')
 313
 314         new_brush.texture = vdm_texture
 315         new_brush.texture_slot.map_mode = 'AREA_PLANE'
 316         new_brush.stroke_method = 'ANCHORED'
 317         new_brush.name = new_brush_name
 318         new_brush.use_color_as_displacement = True
 319         new_brush.strength = 1.0
 320         new_brush.hardness = 0.9
 321
 322         bpy.ops.object.mode_set(mode = default_mode)
 323
 324         if bpy.context.object.mode == 'SCULPT':
 325             context.tool_settings.sculpt.brush = new_brush
 326
 327         return {'FINISHED'}
 328
 329
 330 registered_classes = [
 331     PT_VDMBaker,
 332     vdm_brush_baker_addon_data,
 333     create_vdm_brush,
 334     create_sculpt_plane
 335 ]
 336
 337
 338 def register():
 339     for registered_class in registered_classes:
 340         bpy.utils.register_class(registered_class)
 341
 342     bpy.types.Scene.VDMBrushBakerAddonData = bpy.props.PointerProperty(
 343         type=vdm_brush_baker_addon_data)
 344
 345
 346 def unregister():
 347     for registered_class in registered_classes:
 348         bpy.utils.unregister_class(registered_class)
 349
 350     del bpy.types.Scene.VDMBrushBakerAddonData