refactor

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)

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)))
+        layout.label(text="Estimated ID-Count: " + str(source.estimate_ids(context, props)))

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 []

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"

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

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

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)