From 6a891687070ac3fedade29e2365aa0eb80c0e6fe Mon Sep 17 00:00:00 2001 From: Michel Fedde <35878897+Neintonine@users.noreply.github.com> Date: Tue, 30 Jan 2024 15:51:46 +0100 Subject: [PATCH] refactor --- src/operators/bake_to_id_map.py | 86 +++-------------------------- src/panels/panel_info.py | 5 +- src/panels/panel_options.py | 28 +++------- src/properties/bake_to_id.py | 8 +-- src/types/__init__.py | 43 +++++++++++++++ src/types/sources/material_index.py | 41 ++++++++++++++ src/types/targets/vertex_colors.py | 32 +++++++++++ 7 files changed, 140 insertions(+), 103 deletions(-) create mode 100644 src/types/__init__.py create mode 100644 src/types/sources/material_index.py create mode 100644 src/types/targets/vertex_colors.py diff --git a/src/operators/bake_to_id_map.py b/src/operators/bake_to_id_map.py index 378a897..dcd6643 100644 --- a/src/operators/bake_to_id_map.py +++ b/src/operators/bake_to_id_map.py @@ -3,6 +3,8 @@ import math import bpy +from .. types import (get_source, get_target) + class BakeToIDMapOperator(bpy.types.Operator): bl_idname = "object.bake_to_id_map" @@ -22,17 +24,10 @@ class BakeToIDMapOperator(bpy.types.Operator): return {'FINISHED'} - @staticmethod - def count_ids(context, props): - count = 0 - if props.source == 'MATERIAL_INDEX': - for obj in context.selected_objects: - count += len(obj.material_slots) - - return count - def paint_id_mask(self, context, props): - targets = self.get_targets(context, props) + source = get_source(props.source) + + targets = source.get_targets(context) if len(targets) < 1: return @@ -53,71 +48,6 @@ class BakeToIDMapOperator(bpy.types.Operator): colors.append(colorsys.hls_to_rgb(h, l, s)) - self.paint_targets(props, targets, colors) - self.after_painting(context, props) - - def get_targets(self, context, props): - if props.source == 'MATERIAL_INDEX': - return self.get_material_targets(context) - - return [] - - def get_material_targets(self, context): - targets = [] - for obj in context.selected_objects: - if not obj.material_slots: continue - - mesh = obj.data - if len(obj.material_slots) < 2: - targets.append((mesh, mesh.polygons)) - continue - - polygonMaterials = {} - for poly in mesh.polygons: - if poly.material_index not in polygonMaterials: - polygonMaterials[poly.material_index] = [] - - polygonMaterials[poly.material_index].append(poly) - - for polygons in polygonMaterials.values(): - targets.append((mesh, polygons)) - - return targets - - def paint_targets(self, props, targets, colors): - if props.target == 'VERTEX_COLORS': - self.paint_targets_vertex_colors(props, targets, colors) - - def paint_targets_vertex_colors(self, props, targets, colors): - sortedTargets = {} - for i in range(len(targets)): - target = targets[i] - obj = target[0] - indecies = target[1] - - if obj not in sortedTargets: - sortedTargets[obj] = [] - - sortedTargets[obj].append((indecies, colors[i])) - - layer_name = props.target_vertex_color_attribute_name - for mesh in sortedTargets: - if layer_name in mesh.vertex_colors: - mesh.vertex_colors.remove(mesh.vertex_colors[layer_name]) - - vertex_color_layer = mesh.vertex_colors.new(name=layer_name) - - for (indecies, color) in sortedTargets[mesh]: - for polygon in indecies: - for idx in polygon.loop_indices: - vertex_color_layer.data[idx].color = (color[0], color[1], color[2], 1.0) - - def after_painting(self, context, props): - if props.source == 'MATERIAL_INDEX': - self.after_painting_source_material_index(props, context) - return - - def after_painting_source_material_index(self, props, context): - if props.source_materials_remove_all: - for obj in context.selected_objects: - obj.data.materials.clear() + target = get_target(props.target) + target.paint_targets(props, targets, colors) + source.after_painting(context, props) \ No newline at end of file diff --git a/src/panels/panel_info.py b/src/panels/panel_info.py index 22dde49..7ac9a0f 100644 --- a/src/panels/panel_info.py +++ b/src/panels/panel_info.py @@ -1,6 +1,7 @@ import bpy from .. operators.bake_to_id_map import BakeToIDMapOperator +from ..types import get_source class BakeToIDInfoPanel(bpy.types.Panel): @@ -16,5 +17,7 @@ class BakeToIDInfoPanel(bpy.types.Panel): props = context.scene.bake_to_id_props + source = get_source(props.source) + layout.label(text="Selected Object-Count: " + str(len(context.selected_objects))) - layout.label(text="Estimated ID-Count: " + str(BakeToIDMapOperator.count_ids(context, props))) \ No newline at end of file + layout.label(text="Estimated ID-Count: " + str(source.estimate_ids(context, props))) \ No newline at end of file diff --git a/src/panels/panel_options.py b/src/panels/panel_options.py index 3b74dbe..2d1445e 100644 --- a/src/panels/panel_options.py +++ b/src/panels/panel_options.py @@ -1,5 +1,9 @@ +import textwrap + import bpy +from src.types import get_source, get_target + class BakeToIDOptionsPanel(bpy.types.Panel): bl_idname = "PANEL.BAKE_TO_ID_MAP_PT_SETTINGS_OPTIONS" @@ -15,31 +19,15 @@ class BakeToIDOptionsPanel(bpy.types.Panel): props = context.scene.bake_to_id_props layout.prop(props, "source") + source = get_source(props.source) source_settings_box = layout.box() - source_settings = self.get_source_settings(props) - for setting in source_settings: + for setting in source.connected_properties: source_settings_box.prop(props, setting) layout.separator() layout.prop(props, "target") + target = get_target(props.target) target_settings_box = layout.box() - target_settings = self.get_target_settings(props) - for setting in target_settings: + for setting in target.connected_properties: target_settings_box.prop(props, setting) - - def get_source_settings(self, props): - if props.source == 'MATERIAL_INDEX': - return [ - 'source_materials_remove_all' - ] - - return [] - - def get_target_settings(self, props): - if props.target == 'VERTEX_COLORS': - return [ - 'target_vertex_color_attribute_name' - ] - - return [] \ No newline at end of file diff --git a/src/properties/bake_to_id.py b/src/properties/bake_to_id.py index b7be6f7..77978eb 100644 --- a/src/properties/bake_to_id.py +++ b/src/properties/bake_to_id.py @@ -1,19 +1,19 @@ from bpy.types import (PropertyGroup) from bpy.props import (EnumProperty, BoolProperty, IntProperty, StringProperty) +from src.types import get_source_enum, get_targets_enum + class BakeToIDProperties(PropertyGroup): source: EnumProperty( - items=[ - ('MATERIAL_INDEX', 'Material Index', "It takes the material index as ID basis", 0) - ], + items=get_source_enum(), name="Source", description="From where should the IDs be taken", default = "MATERIAL_INDEX" ) target: EnumProperty( - items=[('VERTEX_COLORS', 'Vertex Colors', "It bakes the ID to the vertex colors", 0)], + items=get_targets_enum(), name="Target", description="To where should the IDs should be baked to", default="VERTEX_COLORS" diff --git a/src/types/__init__.py b/src/types/__init__.py new file mode 100644 index 0000000..c814b4c --- /dev/null +++ b/src/types/__init__.py @@ -0,0 +1,43 @@ +from . sources import material_index as source_mat_index + +from . targets import vertex_colors as target_vertex_colors + +_sources = [ + source_mat_index +] + +_targets = [ + target_vertex_colors +] + +def get_source(id): + for source in _sources: + if source.source_id == id: + return source + + raise Exception("Source not found: " + id) + +def get_target(id): + for target in _targets: + if target.target_id == id: + return target + + raise Exception("Target not found: " + id) + +def get_source_enum(): + enumList = [] + i = 0 + for source in _sources: + enumList.append((source.source_id, source.name, source.description, i)) + i += 1 + + return enumList + +def get_targets_enum(): + enumList = [] + i = 0 + for target in _targets: + enumList.append((target.target_id, target.name, target.description, i)) + i += 1 + + return enumList diff --git a/src/types/sources/material_index.py b/src/types/sources/material_index.py new file mode 100644 index 0000000..d7b1b5b --- /dev/null +++ b/src/types/sources/material_index.py @@ -0,0 +1,41 @@ +source_id = 'MATERIAL_INDEX' +name = 'Material Index' +description = 'Uses the current material index as basis for ID mask.' + +connected_properties = [ + 'source_materials_remove_all' +] + +def get_targets(context): + targets = [] + for obj in context.selected_objects: + if not obj.material_slots: continue + + mesh = obj.data + if len(obj.material_slots) < 2: + targets.append((mesh, mesh.polygons)) + continue + + polygonMaterials = {} + for poly in mesh.polygons: + if poly.material_index not in polygonMaterials: + polygonMaterials[poly.material_index] = [] + + polygonMaterials[poly.material_index].append(poly) + + for polygons in polygonMaterials.values(): + targets.append((mesh, polygons)) + + return targets + +def after_painting(context, props): + if props.source_materials_remove_all: + for obj in context.selected_objects: + obj.data.materials.clear() + +def estimate_ids(context, props): + count = 0 + for obj in context.selected_objects: + count += len(obj.material_slots) + + return count diff --git a/src/types/targets/vertex_colors.py b/src/types/targets/vertex_colors.py new file mode 100644 index 0000000..6b6340d --- /dev/null +++ b/src/types/targets/vertex_colors.py @@ -0,0 +1,32 @@ +target_id = 'VERTEX_COLORS' +name = 'Vertex Colors' +description = 'Bakes the ID onto the vertex color' + +connected_properties = [ + 'target_vertex_color_attribute_name' +] + + +def paint_targets(props, targets, colors): + sortedTargets = {} + for i in range(len(targets)): + target = targets[i] + obj = target[0] + indecies = target[1] + + if obj not in sortedTargets: + sortedTargets[obj] = [] + + sortedTargets[obj].append((indecies, colors[i])) + + layer_name = props.target_vertex_color_attribute_name + for mesh in sortedTargets: + if layer_name in mesh.vertex_colors: + mesh.vertex_colors.remove(mesh.vertex_colors[layer_name]) + + vertex_color_layer = mesh.vertex_colors.new(name=layer_name) + + for (indecies, color) in sortedTargets[mesh]: + for polygon in indecies: + for idx in polygon.loop_indices: + vertex_color_layer.data[idx].color = (color[0], color[1], color[2], 1.0) \ No newline at end of file