caps, visual mode

__init__.py

@@ -2,147 +2,185 @@
 Armature Mesher - Blender Addon
 Converts a selected armature into an envelope-style mesh using each bone's
 head/tail radius values (exactly like Envelope display mode).
+
+Supports:
+  - Object mode  → rest pose (arm.bones)
+  - Pose mode    → current pose (obj.pose.bones)
+  - Edit mode    → current edit-bone positions
+
+Geometry:
+  - UV-sphere at head and tail (head_radius / tail_radius)
+  - Capped frustum cylinder connecting the two spheres
 """
 
 import bpy
 import bmesh
 import math
-from mathutils import Vector, Matrix, Quaternion
+from mathutils import Vector, Matrix
 
 
 # ---------------------------------------------------------------------------
 # Geometry helpers
 # ---------------------------------------------------------------------------
 
-def make_sphere_bmesh(bm, center: Vector, radius: float, segments: int = 12) -> list:
+def _rotation_matrix_to_axis(target_axis: Vector) -> Matrix:
-    """Add a UV-sphere to an existing BMesh, return the new verts."""
+    """Return a 4x4 rotation matrix that rotates Z → target_axis."""
-    verts = []
+    ax = target_axis.normalized()
-    rings = segments // 2
+    z = Vector((0.0, 0.0, 1.0))
-    for ring in range(rings + 1):
+    cross = z.cross(ax)
-        phi = math.pi * ring / rings          # 0 … π
+    if cross.length < 1e-6:
-        for seg in range(segments):
+        if ax.z > 0:
-            theta = 2 * math.pi * seg / segments
+            return Matrix.Identity(4)
-            x = radius * math.sin(phi) * math.cos(theta)
+        else:
-            y = radius * math.sin(phi) * math.sin(theta)
+            return Matrix.Rotation(math.pi, 4, 'X')
-            z = radius * math.cos(phi)
+    angle = math.acos(max(-1.0, min(1.0, z.dot(ax))))
-            verts.append(bm.verts.new(center + Vector((x, y, z))))
+    return Matrix.Rotation(angle, 4, cross.normalized())
 
-    # Connect rings
+
-    for ring in range(rings):
+def add_sphere(bm: bmesh.types.BMesh, center: Vector, radius: float, segments: int):
-        for seg in range(segments):
+    """Add a closed UV-sphere to bm."""
-            nxt = (seg + 1) % segments
+    rings = max(segments // 2, 2)
-            v0 = verts[ring * segments + seg]
+    verts = []
-            v1 = verts[ring * segments + nxt]
+
-            v2 = verts[(ring + 1) * segments + nxt]
+    for r in range(rings + 1):
-            v3 = verts[(ring + 1) * segments + seg]
+        phi = math.pi * r / rings
+        sin_phi = math.sin(phi)
+        cos_phi = math.cos(phi)
+        for s in range(segments):
+            theta = 2.0 * math.pi * s / segments
+            v = bm.verts.new(center + Vector((
+                radius * sin_phi * math.cos(theta),
+                radius * sin_phi * math.sin(theta),
+                radius * cos_phi,
+            )))
+            verts.append(v)
+
+    for r in range(rings):
+        for s in range(segments):
+            s_next = (s + 1) % segments
+            v0 = verts[r       * segments + s]
+            v1 = verts[r       * segments + s_next]
+            v2 = verts[(r + 1) * segments + s_next]
+            v3 = verts[(r + 1) * segments + s]
             try:
                 bm.faces.new((v0, v1, v2, v3))
             except ValueError:
                 pass
+
     return verts
 
 
+def add_capped_frustum(bm: bmesh.types.BMesh,
+                        head_center: Vector, head_radius: float,
+                        tail_center: Vector, tail_radius: float,
+                        segments: int):
+    """
+    Add a capped tapered cylinder between two sphere centers.
+    Rings are placed at the sphere equators (offset = radius along axis)
+    so they sit flush with the sphere surfaces.
+    Each open end is closed with a triangle fan cap.
+    """
+    axis = tail_center - head_center
+    length = axis.length
+    if length < 1e-6:
+        return
 
+    ax = axis.normalized()
+    rot = _rotation_matrix_to_axis(ax)
+
+    # Place rings at the sphere equators along the bone axis
+    head_ring_center = head_center + ax * head_radius
+    tail_ring_center = tail_center - ax * tail_radius
+
+    # Skip if rings would overlap (bone too short relative to radii)
+    if (tail_ring_center - head_ring_center).dot(ax) < 1e-6:
+        return
+
+    # Build the two rings
+    head_ring = []
+    tail_ring = []
+    for i in range(segments):
+        theta = 2.0 * math.pi * i / segments
+        local = Vector((math.cos(theta), math.sin(theta), 0.0))
+        # rot is 4x4; multiply then drop w component
+        offset = (rot @ local.to_4d()).to_3d()
+        head_ring.append(bm.verts.new(head_ring_center + offset * head_radius))
+        tail_ring.append(bm.verts.new(tail_ring_center + offset * tail_radius))
+
+    # Side quads
+    for i in range(segments):
+        nxt = (i + 1) % segments
+        try:
+            bm.faces.new((head_ring[i], head_ring[nxt],
+                          tail_ring[nxt], tail_ring[i]))
+        except ValueError:
+            pass
+
+    # Head cap — fan, winding faces inward (away from tail)
+    head_cap_v = bm.verts.new(head_ring_center)
+    for i in range(segments):
+        nxt = (i + 1) % segments
+        try:
+            bm.faces.new((head_cap_v, head_ring[nxt], head_ring[i]))
+        except ValueError:
+            pass
+
+    # Tail cap — fan, winding faces outward (away from head)
+    tail_cap_v = bm.verts.new(tail_ring_center)
+    for i in range(segments):
+        nxt = (i + 1) % segments
+        try:
+            bm.faces.new((tail_cap_v, tail_ring[i], tail_ring[nxt]))
+        except ValueError:
+            pass
 
 
 # ---------------------------------------------------------------------------
-# Core build function
+# Bone data collection — mode-aware
 # ---------------------------------------------------------------------------
 
-SEGMENTS = 16   # quality — increase for smoother result
+def collect_bone_data(arm_obj) -> list:
+    """
+    Returns a list of dicts with world-space head/tail and radii.
+    Reads from the appropriate source based on the armature's current mode:
+      OBJECT → arm.bones           (rest pose)
+      POSE   → obj.pose.bones      (current posed positions)
+      EDIT   → arm.edit_bones      (what you see in edit mode)
+    """
+    mode      = arm_obj.mode
+    world_mat = arm_obj.matrix_world
+    result    = []
 
+    if mode == 'POSE':
+        for pb in arm_obj.pose.bones:
+            b = pb.bone  # underlying data bone for radii
+            result.append({
+                'head':   world_mat @ pb.head,
+                'tail':   world_mat @ pb.tail,
+                'head_r': max(b.head_radius, 0.001),
+                'tail_r': max(b.tail_radius, 0.001),
+            })
 
-def build_envelope_mesh(armature_obj: bpy.types.Object) -> bpy.types.Object:
+    elif mode == 'EDIT':
-    arm = armature_obj.data
+        # edit_bones is accessible while in edit mode
-    world_mat = armature_obj.matrix_world
+        for eb in arm_obj.data.edit_bones:
+            result.append({
+                'head':   world_mat @ eb.head,
+                'tail':   world_mat @ eb.tail,
+                'head_r': max(eb.head_radius, 0.001),
+                'tail_r': max(eb.tail_radius, 0.001),
+            })
 
-    bm = bmesh.new()
+    else:  # OBJECT / everything else → rest pose
+        for bone in arm_obj.data.bones:
+            result.append({
+                'head':   world_mat @ bone.head_local,
+                'tail':   world_mat @ bone.tail_local,
+                'head_r': max(bone.head_radius, 0.001),
+                'tail_r': max(bone.tail_radius, 0.001),
+            })
 
-    for bone in arm.bones:
+    return result
-        # Bone head/tail in armature local space
-        head_local = bone.head_local          # Vector
-        tail_local = bone.tail_local
-
-        # Convert to world space
-        head_w = world_mat @ head_local
-        tail_w = world_mat @ tail_local
-
-        head_r = bone.head_radius
-        tail_r = bone.tail_radius
-
-        # Clamp radii to something visible
-        head_r = max(head_r, 0.001)
-        tail_r = max(tail_r, 0.001)
-
-        # Head sphere
-        make_sphere_bmesh(bm, head_w, head_r, SEGMENTS)
-
-        # Tail sphere
-        make_sphere_bmesh(bm, tail_w, tail_r, SEGMENTS)
-
-        # Connecting frustum
-        axis = tail_w - head_w
-        length = axis.length
-        if length < 1e-6:
-            continue
-
-        ax = axis.normalized()
-
-        def ring_offset(sphere_r, ring_r):
-            rr = min(ring_r, sphere_r)
-            return math.sqrt(max(sphere_r ** 2 - rr ** 2, 0.0))
-
-        head_off = ring_offset(head_r, head_r)
-        tail_off = ring_offset(tail_r, tail_r)
-
-        ring_head_pt = head_w + ax * head_off
-        ring_tail_pt = tail_w - ax * tail_off
-
-        # Only draw frustum if ring planes don't overlap
-        if (ring_tail_pt - ring_head_pt).dot(ax) < 1e-6:
-            continue
-
-        # Build local rotation matrix aligning Z to bone axis
-        z = Vector((0, 0, 1))
-        cross = z.cross(ax)
-        if cross.length < 1e-6:
-            rot = Matrix.Identity(4) if ax.z > 0 else Matrix.Rotation(math.pi, 4, 'X')
-        else:
-            angle = math.acos(max(-1.0, min(1.0, z.dot(ax))))
-            rot = Matrix.Rotation(angle, 4, cross.normalized())
-
-        head_ring = []
-        tail_ring = []
-        for i in range(SEGMENTS):
-            theta = 2 * math.pi * i / SEGMENTS
-            lv = Vector((math.cos(theta), math.sin(theta), 0.0))
-            rv = rot @ lv
-            head_ring.append(bm.verts.new(ring_head_pt + rv * head_r))
-            tail_ring.append(bm.verts.new(ring_tail_pt + rv * tail_r))
-
-        for i in range(SEGMENTS):
-            nxt = (i + 1) % SEGMENTS
-            try:
-                bm.faces.new((head_ring[i], head_ring[nxt],
-                              tail_ring[nxt], tail_ring[i]))
-            except ValueError:
-                pass
-
-    # Merge overlapping verts (where spheres from adjacent bones touch)
-    bmesh.ops.remove_doubles(bm, verts=bm.verts, dist=0.0001)
-
-    # Create mesh data-block
-    mesh = bpy.data.meshes.new(armature_obj.name + "_envelope_mesh")
-    bm.to_mesh(mesh)
-    bm.free()
-
-    mesh.update()
-
-    # Create object
-    obj = bpy.data.objects.new(armature_obj.name + "_envelope", mesh)
-    bpy.context.collection.objects.link(obj)
-
-    return obj
 
 
 # ---------------------------------------------------------------------------
@@ -150,20 +188,18 @@ def build_envelope_mesh(armature_obj: bpy.types.Object) -> bpy.types.Object:
 # ---------------------------------------------------------------------------
 
 class ARMATURE_OT_build_envelope_mesh(bpy.types.Operator):
-    bl_idname = "armature.build_envelope_mesh"
+    bl_idname  = "armature.build_envelope_mesh"
-    bl_label = "Build Envelope Mesh"
+    bl_label   = "Build Envelope Mesh"
     bl_description = (
-        "Convert the selected armature into a mesh that mirrors its "
+        "Convert the selected armature into an envelope-style mesh. "
-        "Envelope display (spheres at joints, tapered cylinders along bones)"
+        "Object mode = rest pose | Pose mode = current pose | Edit mode = edit bones."
     )
     bl_options = {'REGISTER', 'UNDO'}
 
     segments: bpy.props.IntProperty(
         name="Segments",
-        description="Cylinder / sphere resolution",
+        description="Sphere / cylinder resolution",
-        default=16,
+        default=16, min=4, max=64,
-        min=4,
-        max=64,
     )
 
     @classmethod
@@ -172,23 +208,59 @@ class ARMATURE_OT_build_envelope_mesh(bpy.types.Operator):
         return obj is not None and obj.type == 'ARMATURE'
 
     def execute(self, context):
-        global SEGMENTS
+        arm_obj       = context.active_object
-        SEGMENTS = self.segments
+        original_mode = arm_obj.mode
 
-        arm_obj = context.active_object
+        # Snapshot bone data NOW, while we're still in the original mode.
+        # This is especially important for EDIT mode where edit_bones are live.
+        bones_data = collect_bone_data(arm_obj)
 
-        # Ensure we have edit-mode bone data (rest-pose local coords)
+        # Must be in OBJECT mode to create and link new mesh objects.
-        if arm_obj.mode != 'OBJECT':
+        if original_mode != 'OBJECT':
             bpy.ops.object.mode_set(mode='OBJECT')
 
-        result_obj = build_envelope_mesh(arm_obj)
+        # Build the BMesh
+        bm   = bmesh.new()
+        segs = self.segments
+
+        for bd in bones_data:
+            add_sphere(bm, bd['head'], bd['head_r'], segs)
+            add_sphere(bm, bd['tail'], bd['tail_r'], segs)
+            add_capped_frustum(bm,
+                               bd['head'], bd['head_r'],
+                               bd['tail'], bd['tail_r'],
+                               segs)
+
+        # Merge verts from shared joints (parent/child bone connections)
+        bmesh.ops.remove_doubles(bm, verts=bm.verts, dist=0.0001)
+
+        mesh = bpy.data.meshes.new(arm_obj.name + "_envelope_mesh")
+        bm.to_mesh(mesh)
+        bm.free()
+        mesh.update()
+
+        result_obj = bpy.data.objects.new(arm_obj.name + "_envelope", mesh)
+        context.collection.objects.link(result_obj)
+
+        # Re-enter edit mode on the armature if that's where we came from,
+        # then make the new mesh the active/selected object.
+        if original_mode == 'EDIT':
+            arm_obj.select_set(True)
+            context.view_layer.objects.active = arm_obj
+            bpy.ops.object.mode_set(mode='EDIT')
+            arm_obj.select_set(False)
 
-        # Select the new mesh
         bpy.ops.object.select_all(action='DESELECT')
         result_obj.select_set(True)
         context.view_layer.objects.active = result_obj
 
-        self.report({'INFO'}, f"Created: {result_obj.name}")
+        source_label = {
+            'OBJECT': "rest pose",
+            'POSE':   "current pose",
+            'EDIT':   "edit-bone layout",
+        }.get(original_mode, original_mode)
+
+        self.report({'INFO'}, f"Created '{result_obj.name}' from {source_label}.")
         return {'FINISHED'}
 
 
@@ -197,20 +269,27 @@ class ARMATURE_OT_build_envelope_mesh(bpy.types.Operator):
 # ---------------------------------------------------------------------------
 
 class VIEW3D_PT_armature_mesher(bpy.types.Panel):
-    bl_label = "Armature Mesher"
+    bl_label       = "Armature Mesher"
-    bl_idname = "VIEW3D_PT_armature_mesher"
+    bl_idname      = "VIEW3D_PT_armature_mesher"
-    bl_space_type = 'VIEW_3D'
+    bl_space_type  = 'VIEW_3D'
     bl_region_type = 'UI'
-    bl_category = "Armature"
+    bl_category    = "Armature"
 
     @classmethod
     def poll(cls, context):
-        return (context.active_object is not None and
+        obj = context.active_object
-                context.active_object.type == 'ARMATURE')
+        return obj is not None and obj.type == 'ARMATURE'
 
     def draw(self, context):
         layout = self.layout
-        layout.label(text="Selected: " + context.active_object.name)
+        obj    = context.active_object
+        mode_labels = {
+            'OBJECT': "Object → rest pose",
+            'POSE':   "Pose → current pose",
+            'EDIT':   "Edit → edit bones",
+        }
+        layout.label(text=obj.name)
+        layout.label(text=mode_labels.get(obj.mode, obj.mode), icon='INFO')
         layout.separator()
         op = layout.operator(
             ARMATURE_OT_build_envelope_mesh.bl_idname,
@@ -229,16 +308,13 @@ classes = (
     VIEW3D_PT_armature_mesher,
 )
 
-
 def register():
     for cls in classes:
         bpy.utils.register_class(cls)
 
-
 def unregister():
     for cls in reversed(classes):
         bpy.utils.unregister_class(cls)
 
-
 if __name__ == "__main__":
     register()