907 lines
33 KiB
Python
907 lines
33 KiB
Python
"""
|
|
GLB Brace Generator with Markers
|
|
|
|
This module generates GLB brace files with embedded markers for editing.
|
|
Supports both regular (fitted) and vase-shaped templates.
|
|
|
|
PRESSURE ZONES EXPLANATION:
|
|
===========================
|
|
|
|
The brace has 4 main pressure/expansion zones that correct spinal curvature:
|
|
|
|
1. THORACIC PAD (LM_PAD_TH) - PUSH ZONE
|
|
- Location: On the CONVEX side of the thoracic curve (the side that bulges out)
|
|
- Function: Pushes INWARD to correct the thoracic curvature
|
|
- For right thoracic curves: pad is on the RIGHT back
|
|
- Depth: 8-25mm depending on Cobb angle severity
|
|
|
|
2. THORACIC BAY (LM_BAY_TH) - EXPANSION ZONE
|
|
- Location: OPPOSITE the thoracic pad (concave side)
|
|
- Function: Creates SPACE for the body to move INTO during correction
|
|
- The ribs/body shift into this space as the pad pushes
|
|
- Clearance: 10-35mm
|
|
|
|
3. LUMBAR PAD (LM_PAD_LUM) - PUSH ZONE
|
|
- Location: On the CONVEX side of the lumbar curve
|
|
- Function: Pushes INWARD to correct lumbar curvature
|
|
- Usually on the opposite side of thoracic pad (for S-curves)
|
|
- Depth: 6-20mm
|
|
|
|
4. LUMBAR BAY (LM_BAY_LUM) - EXPANSION ZONE
|
|
- Location: OPPOSITE the lumbar pad
|
|
- Function: Creates SPACE for lumbar correction
|
|
- Clearance: 8-25mm
|
|
|
|
5. HIP ANCHORS (LM_ANCHOR_HIP_L/R) - STABILITY ZONES
|
|
- Location: Around the hip/pelvis area on both sides
|
|
- Function: Grip the pelvis to prevent brace from riding up
|
|
- Slight inward pressure to anchor the brace
|
|
|
|
The Rigo classification determines which zones are primary:
|
|
- A types (3-curve): Strong thoracic pad, minor lumbar
|
|
- B types (4-curve): Both thoracic and lumbar pads are primary
|
|
- C types (non-3-non-4): Balanced thoracic, neutral pelvis
|
|
- E types (single lumbar/TL): Strong lumbar/TL pad, counter-thoracic
|
|
|
|
"""
|
|
|
|
import json
|
|
import numpy as np
|
|
import trimesh
|
|
from pathlib import Path
|
|
from typing import Dict, Any, Optional, Tuple, Literal
|
|
from dataclasses import dataclass, asdict
|
|
|
|
# Paths to template directories
|
|
BASE_DIR = Path(__file__).parent.parent
|
|
BRACES_DIR = BASE_DIR / "braces"
|
|
REGULAR_TEMPLATES_DIR = BRACES_DIR / "brace_templates"
|
|
VASE_TEMPLATES_DIR = BRACES_DIR / "vase_brace_templates"
|
|
|
|
# Template types
|
|
TemplateType = Literal["regular", "vase"]
|
|
|
|
@dataclass
|
|
class MarkerPositions:
|
|
"""Marker positions for a brace template."""
|
|
LM_PELVIS_CENTER: Tuple[float, float, float]
|
|
LM_TOP_CENTER: Tuple[float, float, float]
|
|
LM_PAD_TH: Tuple[float, float, float]
|
|
LM_BAY_TH: Tuple[float, float, float]
|
|
LM_PAD_LUM: Tuple[float, float, float]
|
|
LM_BAY_LUM: Tuple[float, float, float]
|
|
LM_ANCHOR_HIP_L: Tuple[float, float, float]
|
|
LM_ANCHOR_HIP_R: Tuple[float, float, float]
|
|
|
|
@dataclass
|
|
class PressureZone:
|
|
"""Describes a pressure or expansion zone on the brace."""
|
|
name: str
|
|
marker_name: str
|
|
position: Tuple[float, float, float]
|
|
zone_type: Literal["pad", "bay", "anchor"]
|
|
function: str
|
|
direction: Literal["inward", "outward", "grip"]
|
|
depth_mm: float = 0.0
|
|
radius_mm: Tuple[float, float, float] = (50.0, 80.0, 40.0)
|
|
|
|
@dataclass
|
|
class BraceGenerationResult:
|
|
"""Result of brace generation with markers."""
|
|
glb_path: str
|
|
stl_path: str
|
|
json_path: str
|
|
template_type: str
|
|
rigo_type: str
|
|
markers: Dict[str, Tuple[float, float, float]]
|
|
basis: Dict[str, Any]
|
|
pressure_zones: list
|
|
mesh_stats: Dict[str, int]
|
|
transform_applied: Optional[Dict[str, Any]] = None
|
|
|
|
|
|
def get_template_paths(rigo_type: str, template_type: TemplateType) -> Tuple[Path, Path]:
|
|
"""
|
|
Get paths to GLB template and markers JSON.
|
|
|
|
Args:
|
|
rigo_type: Rigo classification (A1, A2, A3, B1, B2, C1, C2, E1, E2)
|
|
template_type: "regular" or "vase"
|
|
|
|
Returns:
|
|
Tuple of (glb_path, markers_json_path)
|
|
"""
|
|
if template_type == "regular":
|
|
glb_path = REGULAR_TEMPLATES_DIR / f"{rigo_type}_marked_v3.glb"
|
|
json_path = REGULAR_TEMPLATES_DIR / f"{rigo_type}_marked_v3.markers.json"
|
|
else: # vase
|
|
glb_path = VASE_TEMPLATES_DIR / "glb" / f"{rigo_type}_vase_marked.glb"
|
|
json_path = VASE_TEMPLATES_DIR / "markers_json" / f"{rigo_type}_vase_marked.markers.json"
|
|
|
|
return glb_path, json_path
|
|
|
|
|
|
def load_template_markers(rigo_type: str, template_type: TemplateType) -> Dict[str, Any]:
|
|
"""Load markers from JSON file for a template."""
|
|
_, json_path = get_template_paths(rigo_type, template_type)
|
|
|
|
if not json_path.exists():
|
|
raise FileNotFoundError(f"Markers JSON not found: {json_path}")
|
|
|
|
with open(json_path, "r") as f:
|
|
return json.load(f)
|
|
|
|
|
|
def load_glb_template(rigo_type: str, template_type: TemplateType) -> trimesh.Trimesh:
|
|
"""Load GLB template as trimesh."""
|
|
glb_path, _ = get_template_paths(rigo_type, template_type)
|
|
|
|
if not glb_path.exists():
|
|
raise FileNotFoundError(f"GLB template not found: {glb_path}")
|
|
|
|
scene = trimesh.load(str(glb_path))
|
|
|
|
# If it's a scene, concatenate all meshes
|
|
if isinstance(scene, trimesh.Scene):
|
|
meshes = [g for g in scene.geometry.values() if isinstance(g, trimesh.Trimesh)]
|
|
if meshes:
|
|
mesh = trimesh.util.concatenate(meshes)
|
|
else:
|
|
raise ValueError(f"No valid meshes found in GLB: {glb_path}")
|
|
else:
|
|
mesh = scene
|
|
|
|
return mesh
|
|
|
|
|
|
def calculate_pressure_zones(
|
|
markers: Dict[str, Any],
|
|
rigo_type: str,
|
|
cobb_angles: Dict[str, float]
|
|
) -> list:
|
|
"""
|
|
Calculate pressure zone parameters based on markers and analysis.
|
|
|
|
Args:
|
|
markers: Marker positions from template
|
|
rigo_type: Rigo classification
|
|
cobb_angles: Dict with PT, MT, TL Cobb angles
|
|
|
|
Returns:
|
|
List of PressureZone objects
|
|
"""
|
|
marker_pos = markers.get("markers", markers)
|
|
|
|
# Calculate severity from Cobb angles
|
|
mt_angle = cobb_angles.get("MT", 0)
|
|
tl_angle = cobb_angles.get("TL", 0)
|
|
|
|
# Severity mapping: Cobb -> depth
|
|
def cobb_to_depth(angle: float, min_depth: float = 6.0, max_depth: float = 22.0) -> float:
|
|
severity = min(max((angle - 10) / 40, 0), 1) # 0-1 range
|
|
return min_depth + severity * (max_depth - min_depth)
|
|
|
|
th_depth = cobb_to_depth(mt_angle, 8.0, 22.0)
|
|
lum_depth = cobb_to_depth(tl_angle, 6.0, 18.0)
|
|
|
|
# Bay clearance is typically 1.2-1.5x pad depth
|
|
th_clearance = th_depth * 1.3 + 5
|
|
lum_clearance = lum_depth * 1.3 + 4
|
|
|
|
zones = [
|
|
PressureZone(
|
|
name="Thoracic Pad",
|
|
marker_name="LM_PAD_TH",
|
|
position=tuple(marker_pos.get("LM_PAD_TH", [0, 0, 0])),
|
|
zone_type="pad",
|
|
function="Pushes INWARD on thoracic curve convex side to correct curvature",
|
|
direction="inward",
|
|
depth_mm=th_depth,
|
|
radius_mm=(50.0, 90.0, 40.0)
|
|
),
|
|
PressureZone(
|
|
name="Thoracic Bay",
|
|
marker_name="LM_BAY_TH",
|
|
position=tuple(marker_pos.get("LM_BAY_TH", [0, 0, 0])),
|
|
zone_type="bay",
|
|
function="Creates SPACE on thoracic concave side for body to shift into",
|
|
direction="outward",
|
|
depth_mm=th_clearance,
|
|
radius_mm=(65.0, 110.0, 55.0)
|
|
),
|
|
PressureZone(
|
|
name="Lumbar Pad",
|
|
marker_name="LM_PAD_LUM",
|
|
position=tuple(marker_pos.get("LM_PAD_LUM", [0, 0, 0])),
|
|
zone_type="pad",
|
|
function="Pushes INWARD on lumbar curve convex side to correct curvature",
|
|
direction="inward",
|
|
depth_mm=lum_depth,
|
|
radius_mm=(55.0, 85.0, 45.0)
|
|
),
|
|
PressureZone(
|
|
name="Lumbar Bay",
|
|
marker_name="LM_BAY_LUM",
|
|
position=tuple(marker_pos.get("LM_BAY_LUM", [0, 0, 0])),
|
|
zone_type="bay",
|
|
function="Creates SPACE on lumbar concave side for body to shift into",
|
|
direction="outward",
|
|
depth_mm=lum_clearance,
|
|
radius_mm=(70.0, 100.0, 60.0)
|
|
),
|
|
PressureZone(
|
|
name="Left Hip Anchor",
|
|
marker_name="LM_ANCHOR_HIP_L",
|
|
position=tuple(marker_pos.get("LM_ANCHOR_HIP_L", [0, 0, 0])),
|
|
zone_type="anchor",
|
|
function="Grips left hip/pelvis to stabilize brace and prevent riding up",
|
|
direction="grip",
|
|
depth_mm=4.0,
|
|
radius_mm=(40.0, 60.0, 40.0)
|
|
),
|
|
PressureZone(
|
|
name="Right Hip Anchor",
|
|
marker_name="LM_ANCHOR_HIP_R",
|
|
position=tuple(marker_pos.get("LM_ANCHOR_HIP_R", [0, 0, 0])),
|
|
zone_type="anchor",
|
|
function="Grips right hip/pelvis to stabilize brace and prevent riding up",
|
|
direction="grip",
|
|
depth_mm=4.0,
|
|
radius_mm=(40.0, 60.0, 40.0)
|
|
),
|
|
]
|
|
|
|
return zones
|
|
|
|
|
|
def transform_markers(
|
|
markers: Dict[str, list],
|
|
transform_matrix: np.ndarray
|
|
) -> Dict[str, Tuple[float, float, float]]:
|
|
"""Apply transformation matrix to all marker positions."""
|
|
transformed = {}
|
|
|
|
for name, pos in markers.items():
|
|
if isinstance(pos, (list, tuple)) and len(pos) == 3:
|
|
# Convert to homogeneous coordinates
|
|
pos_h = np.array([pos[0], pos[1], pos[2], 1.0])
|
|
# Apply transform
|
|
new_pos = transform_matrix @ pos_h
|
|
transformed[name] = (float(new_pos[0]), float(new_pos[1]), float(new_pos[2]))
|
|
|
|
return transformed
|
|
|
|
|
|
def generate_glb_brace(
|
|
rigo_type: str,
|
|
template_type: TemplateType,
|
|
output_dir: Path,
|
|
case_id: str,
|
|
cobb_angles: Dict[str, float],
|
|
body_scan_path: Optional[str] = None,
|
|
clearance_mm: float = 8.0
|
|
) -> BraceGenerationResult:
|
|
"""
|
|
Generate a GLB brace with markers.
|
|
|
|
Args:
|
|
rigo_type: Rigo classification (A1, A2, etc.)
|
|
template_type: "regular" or "vase"
|
|
output_dir: Directory for output files
|
|
case_id: Case identifier
|
|
cobb_angles: Dict with PT, MT, TL angles
|
|
body_scan_path: Optional path to body scan STL for fitting
|
|
clearance_mm: Clearance between body and brace
|
|
|
|
Returns:
|
|
BraceGenerationResult with paths and marker info
|
|
"""
|
|
output_dir = Path(output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Load template
|
|
mesh = load_glb_template(rigo_type, template_type)
|
|
marker_data = load_template_markers(rigo_type, template_type)
|
|
|
|
markers = marker_data.get("markers", {})
|
|
basis = marker_data.get("basis", {})
|
|
|
|
transform_matrix = np.eye(4)
|
|
transform_info = None
|
|
|
|
# If body scan provided, fit to body
|
|
if body_scan_path and Path(body_scan_path).exists():
|
|
mesh, transform_matrix, transform_info = fit_brace_to_body(
|
|
mesh, body_scan_path, clearance_mm, brace_basis=basis, template_type=template_type,
|
|
markers=markers # Pass markers for zone-aware ironing
|
|
)
|
|
# Transform markers to match the body-fitted mesh
|
|
markers = transform_markers(markers, transform_matrix)
|
|
|
|
# Calculate pressure zones
|
|
pressure_zones = calculate_pressure_zones(markers, rigo_type, cobb_angles)
|
|
|
|
# Output file names
|
|
type_suffix = "vase" if template_type == "vase" else "regular"
|
|
glb_filename = f"{case_id}_{rigo_type}_{type_suffix}.glb"
|
|
stl_filename = f"{case_id}_{rigo_type}_{type_suffix}.stl"
|
|
json_filename = f"{case_id}_{rigo_type}_{type_suffix}_markers.json"
|
|
|
|
glb_path = output_dir / glb_filename
|
|
stl_path = output_dir / stl_filename
|
|
json_path = output_dir / json_filename
|
|
|
|
# Export GLB
|
|
mesh.export(str(glb_path))
|
|
|
|
# Export STL
|
|
mesh.export(str(stl_path))
|
|
|
|
# Build output JSON with markers and zones
|
|
output_data = {
|
|
"case_id": case_id,
|
|
"rigo_type": rigo_type,
|
|
"template_type": template_type,
|
|
"cobb_angles": cobb_angles,
|
|
"markers": {k: list(v) if isinstance(v, tuple) else v for k, v in markers.items()},
|
|
"basis": basis,
|
|
"pressure_zones": [
|
|
{
|
|
"name": z.name,
|
|
"marker_name": z.marker_name,
|
|
"position": list(z.position),
|
|
"zone_type": z.zone_type,
|
|
"function": z.function,
|
|
"direction": z.direction,
|
|
"depth_mm": z.depth_mm,
|
|
"radius_mm": list(z.radius_mm)
|
|
}
|
|
for z in pressure_zones
|
|
],
|
|
"mesh_stats": {
|
|
"vertices": len(mesh.vertices),
|
|
"faces": len(mesh.faces)
|
|
},
|
|
"outputs": {
|
|
"glb": str(glb_path),
|
|
"stl": str(stl_path),
|
|
"json": str(json_path)
|
|
}
|
|
}
|
|
|
|
if transform_info:
|
|
output_data["body_fitting"] = transform_info
|
|
|
|
# Save JSON
|
|
with open(json_path, "w") as f:
|
|
json.dump(output_data, f, indent=2)
|
|
|
|
return BraceGenerationResult(
|
|
glb_path=str(glb_path),
|
|
stl_path=str(stl_path),
|
|
json_path=str(json_path),
|
|
template_type=template_type,
|
|
rigo_type=rigo_type,
|
|
markers=markers,
|
|
basis=basis,
|
|
pressure_zones=[asdict(z) for z in pressure_zones],
|
|
mesh_stats=output_data["mesh_stats"],
|
|
transform_applied=transform_info
|
|
)
|
|
|
|
|
|
def iron_brace_to_body(
|
|
brace_mesh: trimesh.Trimesh,
|
|
body_mesh: trimesh.Trimesh,
|
|
min_clearance_mm: float = 3.0,
|
|
max_clearance_mm: float = 15.0,
|
|
smoothing_iterations: int = 2,
|
|
up_axis: int = 2,
|
|
markers: Optional[Dict[str, Any]] = None
|
|
) -> trimesh.Trimesh:
|
|
"""
|
|
Iron the brace surface to conform to the body scan surface.
|
|
|
|
This ensures the brace follows the body contour without excessive gaps.
|
|
Uses zone-aware ironing:
|
|
- FRONT (belly) and BACK: Aggressive ironing for tight fit
|
|
- SIDES (where pads/bays are): Preserve correction zones, moderate ironing
|
|
|
|
Args:
|
|
brace_mesh: The brace mesh to iron
|
|
body_mesh: The body scan mesh to conform to
|
|
min_clearance_mm: Minimum distance from body surface
|
|
max_clearance_mm: Maximum distance from body surface (trigger ironing)
|
|
smoothing_iterations: Number of Laplacian smoothing passes after ironing
|
|
up_axis: Which axis is "up" (0=X, 1=Y, 2=Z)
|
|
markers: Optional dict of marker positions to preserve pressure zones
|
|
|
|
Returns:
|
|
Ironed brace mesh
|
|
"""
|
|
from scipy.spatial import cKDTree
|
|
import math
|
|
|
|
print(f"Ironing brace to body surface (clearance: {min_clearance_mm}-{max_clearance_mm}mm)")
|
|
|
|
# Create a copy to modify
|
|
ironed_mesh = brace_mesh.copy()
|
|
vertices = ironed_mesh.vertices.copy()
|
|
|
|
# Get body center and bounds
|
|
body_center = body_mesh.centroid
|
|
body_bounds = body_mesh.bounds
|
|
|
|
# Determine the torso region (process middle 80% of body height)
|
|
body_height = body_bounds[1, up_axis] - body_bounds[0, up_axis]
|
|
torso_bottom = body_bounds[0, up_axis] + body_height * 0.10
|
|
torso_top = body_bounds[0, up_axis] + body_height * 0.90
|
|
|
|
# Build KD-tree from body mesh vertices for fast nearest neighbor queries
|
|
body_tree = cKDTree(body_mesh.vertices)
|
|
|
|
# Find closest points on body for ALL brace vertices at once
|
|
distances, closest_indices = body_tree.query(vertices, k=1)
|
|
closest_points = body_mesh.vertices[closest_indices]
|
|
|
|
# Determine horizontal axes (perpendicular to up axis)
|
|
horiz_axes = [i for i in range(3) if i != up_axis]
|
|
|
|
# Calculate brace center for angle computation
|
|
brace_center = np.mean(vertices, axis=0)
|
|
|
|
# Identify marker exclusion zones (preserve correction areas)
|
|
exclusion_zones = []
|
|
if markers:
|
|
# Pad and bay markers need preservation
|
|
for marker_name in ['LM_PAD_TH', 'LM_PAD_LUM', 'LM_BAY_TH', 'LM_BAY_LUM']:
|
|
if marker_name in markers:
|
|
pos = markers[marker_name]
|
|
if isinstance(pos, (list, tuple)) and len(pos) >= 3:
|
|
exclusion_zones.append({
|
|
'center': np.array(pos),
|
|
'radius': 60.0, # 60mm exclusion radius around markers
|
|
'name': marker_name
|
|
})
|
|
|
|
# Process each brace vertex
|
|
adjusted_count = 0
|
|
pulled_in_count = 0
|
|
pushed_out_count = 0
|
|
skipped_zone_count = 0
|
|
|
|
# Height normalization
|
|
brace_min_z = vertices[:, up_axis].min()
|
|
brace_max_z = vertices[:, up_axis].max()
|
|
brace_height_range = max(brace_max_z - brace_min_z, 1.0)
|
|
|
|
for i in range(len(vertices)):
|
|
vertex = vertices[i]
|
|
closest_pt = closest_points[i]
|
|
dist = distances[i]
|
|
|
|
# Only process vertices in the torso region
|
|
if vertex[up_axis] < torso_bottom or vertex[up_axis] > torso_top:
|
|
continue
|
|
|
|
# Check if vertex is in an exclusion zone (near pad/bay markers)
|
|
in_exclusion = False
|
|
for zone in exclusion_zones:
|
|
zone_dist = np.linalg.norm(vertex - zone['center'])
|
|
if zone_dist < zone['radius']:
|
|
in_exclusion = True
|
|
skipped_zone_count += 1
|
|
break
|
|
|
|
if in_exclusion:
|
|
continue
|
|
|
|
# Calculate angular position around body center (horizontal plane)
|
|
# 0° = front (belly), 90° = right side, 180° = back, 270° = left side
|
|
rel_pos = vertex - body_center
|
|
angle = math.atan2(rel_pos[horiz_axes[1]], rel_pos[horiz_axes[0]])
|
|
angle_deg = math.degrees(angle) % 360
|
|
|
|
# Determine zone based on angle:
|
|
# FRONT (belly): 315-45° - aggressive ironing
|
|
# BACK: 135-225° - aggressive ironing
|
|
# SIDES: 45-135° and 225-315° - moderate ironing (correction zones)
|
|
is_front_back = (angle_deg < 45 or angle_deg > 315) or (135 < angle_deg < 225)
|
|
|
|
# Height-based clearance adjustment
|
|
height_norm = (vertex[up_axis] - brace_min_z) / brace_height_range
|
|
|
|
# Set clearances based on zone
|
|
if is_front_back:
|
|
# FRONT/BACK: Aggressive ironing - very tight fit
|
|
local_min = min_clearance_mm * 0.5 # Allow closer to body
|
|
local_max = max_clearance_mm * 0.6 # Trigger ironing earlier
|
|
local_target = min_clearance_mm + 2.0 # Target just above minimum
|
|
else:
|
|
# SIDES: More conservative - preserve room for correction
|
|
local_min = min_clearance_mm
|
|
local_max = max_clearance_mm * 1.2 # Allow slightly more gap
|
|
local_target = (min_clearance_mm + max_clearance_mm) / 2
|
|
|
|
# Height adjustments (tighter at hips and chest)
|
|
if height_norm < 0.25 or height_norm > 0.75:
|
|
local_max *= 0.8 # Tighter at extremes
|
|
local_target *= 0.85
|
|
|
|
# Direction from body surface to brace vertex
|
|
direction = vertex - closest_pt
|
|
dir_length = np.linalg.norm(direction)
|
|
|
|
if dir_length < 1e-6:
|
|
direction = vertex - body_center
|
|
direction[up_axis] = 0
|
|
dir_length = np.linalg.norm(direction)
|
|
if dir_length < 1e-6:
|
|
continue
|
|
|
|
direction = direction / dir_length
|
|
|
|
# Determine signed distance
|
|
vertex_dist_to_center = np.linalg.norm(vertex[:2] - body_center[:2])
|
|
closest_dist_to_center = np.linalg.norm(closest_pt[:2] - body_center[:2])
|
|
|
|
if vertex_dist_to_center >= closest_dist_to_center:
|
|
signed_distance = dist
|
|
else:
|
|
signed_distance = -dist
|
|
|
|
# Determine if adjustment is needed
|
|
needs_adjustment = False
|
|
new_position = vertex.copy()
|
|
|
|
if signed_distance > local_max:
|
|
# Gap too large - pull vertex closer to body
|
|
new_position = closest_pt + direction * local_target
|
|
new_position[up_axis] = vertex[up_axis] # Preserve height
|
|
needs_adjustment = True
|
|
pulled_in_count += 1
|
|
|
|
elif signed_distance < local_min:
|
|
# Too close or inside body - push outward
|
|
offset = local_min + 1.0
|
|
outward_dir = closest_pt - body_center
|
|
outward_dir[up_axis] = 0
|
|
outward_length = np.linalg.norm(outward_dir)
|
|
if outward_length > 1e-6:
|
|
outward_dir = outward_dir / outward_length
|
|
new_position = closest_pt + outward_dir * offset
|
|
new_position[up_axis] = vertex[up_axis]
|
|
needs_adjustment = True
|
|
pushed_out_count += 1
|
|
|
|
if needs_adjustment:
|
|
vertices[i] = new_position
|
|
adjusted_count += 1
|
|
|
|
print(f"Ironing adjusted {adjusted_count} vertices (pulled in: {pulled_in_count}, pushed out: {pushed_out_count}, skipped zones: {skipped_zone_count})")
|
|
|
|
# Apply modified vertices
|
|
ironed_mesh.vertices = vertices
|
|
|
|
# Apply Laplacian smoothing to blend changes and remove artifacts
|
|
if smoothing_iterations > 0 and adjusted_count > 0:
|
|
print(f"Applying {smoothing_iterations} smoothing iterations")
|
|
try:
|
|
ironed_mesh = trimesh.smoothing.filter_laplacian(
|
|
ironed_mesh,
|
|
lamb=0.3, # Gentler smoothing to preserve shape
|
|
iterations=smoothing_iterations,
|
|
implicit_time_integration=False
|
|
)
|
|
except Exception as e:
|
|
print(f"Smoothing failed (non-critical): {e}")
|
|
|
|
# Ensure mesh is valid
|
|
ironed_mesh.fix_normals()
|
|
|
|
return ironed_mesh
|
|
|
|
|
|
def fit_brace_to_body(
|
|
brace_mesh: trimesh.Trimesh,
|
|
body_scan_path: str,
|
|
clearance_mm: float = 8.0,
|
|
brace_basis: Optional[Dict[str, Any]] = None,
|
|
template_type: str = "regular",
|
|
enable_ironing: bool = True,
|
|
markers: Optional[Dict[str, Any]] = None
|
|
) -> Tuple[trimesh.Trimesh, np.ndarray, Dict[str, Any]]:
|
|
"""
|
|
Fit brace to body scan using basis alignment.
|
|
|
|
The brace needs to be:
|
|
1. Rotated so its UP axis aligns with body's UP axis (typically Z for body scans)
|
|
2. Scaled to fit around the body with proper clearance
|
|
3. Positioned at the torso level
|
|
4. Ironed to conform to body surface (respecting correction zones)
|
|
|
|
Returns:
|
|
Tuple of (transformed_mesh, transform_matrix, fitting_info)
|
|
"""
|
|
# Load body scan
|
|
body_mesh = trimesh.load(body_scan_path, force='mesh')
|
|
|
|
# Get body dimensions
|
|
body_bounds = body_mesh.bounds
|
|
body_extents = body_mesh.extents
|
|
body_center = body_mesh.centroid
|
|
|
|
# Determine body up axis (typically the longest dimension = height)
|
|
# For human body scans, this is usually Z (from 3D scanners) or Y
|
|
body_up_axis_idx = np.argmax(body_extents)
|
|
print(f"Body up axis: {['X', 'Y', 'Z'][body_up_axis_idx]}, extents: {body_extents}")
|
|
|
|
# Get brace dimensions
|
|
brace_bounds = brace_mesh.bounds
|
|
brace_extents = brace_mesh.extents
|
|
brace_center = brace_mesh.centroid
|
|
|
|
print(f"Brace original extents: {brace_extents}, template_type: {template_type}")
|
|
|
|
# Start building transformation
|
|
transformed_mesh = brace_mesh.copy()
|
|
transform = np.eye(4)
|
|
|
|
# Step 1: Center brace at origin
|
|
T_center = np.eye(4)
|
|
T_center[:3, 3] = -brace_center
|
|
transformed_mesh.apply_transform(T_center)
|
|
transform = T_center @ transform
|
|
|
|
# Step 2: Apply rotations based on template type and body orientation
|
|
# Regular templates have: negative Y is up (inverted), need to flip
|
|
# Vase templates have: positive Y is up
|
|
# Body scan is Z-up
|
|
|
|
if body_up_axis_idx == 2: # Body is Z-up (standard for 3D scanners)
|
|
if template_type == "regular":
|
|
# Regular brace: -Y is up (inverted)
|
|
# 1. Rotate -90° around X to bring Y-up to Z-up
|
|
R1 = trimesh.transformations.rotation_matrix(-np.pi/2, [1, 0, 0])
|
|
transformed_mesh.apply_transform(R1)
|
|
transform = R1 @ transform
|
|
|
|
# 2. The brace is now Z-up but inverted (pelvis at top, shoulders at bottom)
|
|
# Flip 180° around X to correct (this keeps Z as up axis)
|
|
R2 = trimesh.transformations.rotation_matrix(np.pi, [1, 0, 0])
|
|
transformed_mesh.apply_transform(R2)
|
|
transform = R2 @ transform
|
|
|
|
# 3. Rotate around Z to face forward correctly
|
|
R3 = trimesh.transformations.rotation_matrix(-np.pi/2, [0, 0, 1])
|
|
transformed_mesh.apply_transform(R3)
|
|
transform = R3 @ transform
|
|
|
|
print(f"Applied regular brace rotations: X-90°, X+180° (flip), Z-90°")
|
|
|
|
else: # vase
|
|
# Vase brace: positive Y is up
|
|
# 1. Rotate -90° around X to bring Y-up to Z-up
|
|
R1 = trimesh.transformations.rotation_matrix(-np.pi/2, [1, 0, 0])
|
|
transformed_mesh.apply_transform(R1)
|
|
transform = R1 @ transform
|
|
|
|
# 2. Flip 180° around Y to correct orientation (right-side up)
|
|
R2 = trimesh.transformations.rotation_matrix(np.pi, [0, 1, 0])
|
|
transformed_mesh.apply_transform(R2)
|
|
transform = R2 @ transform
|
|
|
|
print(f"Applied vase brace rotations: X-90°, Y+180° (flip)")
|
|
|
|
# Step 3: Get new brace dimensions after rotation
|
|
new_brace_extents = transformed_mesh.extents
|
|
new_brace_center = transformed_mesh.centroid
|
|
print(f"Brace extents after rotation: {new_brace_extents}")
|
|
|
|
# Step 4: Calculate NON-UNIFORM scaling based on body dimensions
|
|
# The brace should cover the TORSO region (~50% of body height)
|
|
# AND wrap around the body with proper girth
|
|
|
|
body_height = body_extents[body_up_axis_idx]
|
|
brace_height = new_brace_extents[body_up_axis_idx] # After rotation, this is the height
|
|
|
|
# Body horizontal dimensions (girth at torso level)
|
|
horizontal_axes = [i for i in range(3) if i != body_up_axis_idx]
|
|
body_width = body_extents[horizontal_axes[0]] # X width
|
|
body_depth = body_extents[horizontal_axes[1]] # Y depth
|
|
|
|
# Brace horizontal dimensions
|
|
brace_width = new_brace_extents[horizontal_axes[0]]
|
|
brace_depth = new_brace_extents[horizontal_axes[1]]
|
|
|
|
# Target: brace height should cover ~65% of body height (full torso coverage)
|
|
target_height = body_height * 0.65
|
|
height_scale = target_height / brace_height if brace_height > 0 else 1.0
|
|
|
|
# Target: brace width/depth should be LARGER than body to wrap AROUND it
|
|
# The brace sits OUTSIDE the body, only pressure points push inward
|
|
# Add ~25% extra + clearance so brace externals are visible outside body
|
|
target_width = body_width * 1.25 + clearance_mm * 2
|
|
target_depth = body_depth * 1.25 + clearance_mm * 2
|
|
|
|
width_scale = target_width / brace_width if brace_width > 0 else 1.0
|
|
depth_scale = target_depth / brace_depth if brace_depth > 0 else 1.0
|
|
|
|
# Apply non-uniform scaling
|
|
# Determine which axis is which after rotation
|
|
S = np.eye(4)
|
|
if body_up_axis_idx == 2: # Z is up
|
|
S[0, 0] = width_scale # X scale
|
|
S[1, 1] = depth_scale # Y scale
|
|
S[2, 2] = height_scale # Z scale (height)
|
|
elif body_up_axis_idx == 1: # Y is up
|
|
S[0, 0] = width_scale # X scale
|
|
S[1, 1] = height_scale # Y scale (height)
|
|
S[2, 2] = depth_scale # Z scale
|
|
else: # X is up (unusual)
|
|
S[0, 0] = height_scale # X scale (height)
|
|
S[1, 1] = width_scale # Y scale
|
|
S[2, 2] = depth_scale # Z scale
|
|
|
|
# Limit scales to reasonable range
|
|
S[0, 0] = max(0.5, min(S[0, 0], 50.0))
|
|
S[1, 1] = max(0.5, min(S[1, 1], 50.0))
|
|
S[2, 2] = max(0.5, min(S[2, 2], 50.0))
|
|
|
|
transformed_mesh.apply_transform(S)
|
|
transform = S @ transform
|
|
|
|
print(f"Applied non-uniform scale: width={S[0,0]:.2f}, depth={S[1,1]:.2f}, height={S[2,2]:.2f}")
|
|
print(f"Target dimensions: width={target_width:.1f}, depth={target_depth:.1f}, height={target_height:.1f}")
|
|
|
|
# For fitting_info, use average scale
|
|
scale = (S[0, 0] + S[1, 1] + S[2, 2]) / 3
|
|
|
|
# Step 6: Position brace at torso level
|
|
# Calculate where the torso is (middle portion of body height)
|
|
body_height = body_extents[body_up_axis_idx]
|
|
body_bottom = body_bounds[0, body_up_axis_idx]
|
|
body_top = body_bounds[1, body_up_axis_idx]
|
|
|
|
# Torso is roughly the middle 40% of body height (from ~30% to ~70%)
|
|
torso_center_ratio = 0.5 # Middle of body
|
|
torso_center_height = body_bottom + body_height * torso_center_ratio
|
|
|
|
# Target position: center horizontally on body, at torso height vertically
|
|
target_center = body_center.copy()
|
|
target_center[body_up_axis_idx] = torso_center_height
|
|
|
|
# Current brace center after transformations
|
|
current_center = transformed_mesh.centroid
|
|
|
|
T_position = np.eye(4)
|
|
T_position[:3, 3] = target_center - current_center
|
|
transformed_mesh.apply_transform(T_position)
|
|
transform = T_position @ transform
|
|
|
|
# Step 7: Iron brace to conform to body surface (eliminate gaps and humps)
|
|
# Transform markers so we can exclude correction zones from ironing
|
|
transformed_markers = None
|
|
if markers:
|
|
transformed_markers = transform_markers(markers, transform)
|
|
|
|
ironing_info = {}
|
|
if enable_ironing:
|
|
try:
|
|
print(f"Starting brace ironing to body surface...")
|
|
pre_iron_extents = transformed_mesh.extents.copy()
|
|
|
|
transformed_mesh = iron_brace_to_body(
|
|
brace_mesh=transformed_mesh,
|
|
body_mesh=body_mesh,
|
|
min_clearance_mm=clearance_mm * 0.4, # Allow closer for tight fit
|
|
max_clearance_mm=clearance_mm * 1.5, # Iron areas with gaps > 1.5x clearance
|
|
smoothing_iterations=3,
|
|
up_axis=body_up_axis_idx,
|
|
markers=transformed_markers
|
|
)
|
|
|
|
post_iron_extents = transformed_mesh.extents
|
|
ironing_info = {
|
|
"enabled": True,
|
|
"pre_iron_extents": pre_iron_extents.tolist(),
|
|
"post_iron_extents": post_iron_extents.tolist(),
|
|
"min_clearance_mm": clearance_mm * 0.5,
|
|
"max_clearance_mm": clearance_mm * 2.0,
|
|
}
|
|
print(f"Ironing complete. Extents changed from {pre_iron_extents} to {post_iron_extents}")
|
|
|
|
except Exception as e:
|
|
print(f"Ironing failed (non-critical): {e}")
|
|
ironing_info = {"enabled": False, "error": str(e)}
|
|
else:
|
|
ironing_info = {"enabled": False}
|
|
|
|
fitting_info = {
|
|
"scale_avg": float(scale),
|
|
"scale_x": float(S[0, 0]),
|
|
"scale_y": float(S[1, 1]),
|
|
"scale_z": float(S[2, 2]),
|
|
"template_type": template_type,
|
|
"body_extents": body_extents.tolist(),
|
|
"brace_extents_original": brace_extents.tolist(),
|
|
"brace_extents_final": transformed_mesh.extents.tolist(),
|
|
"clearance_mm": clearance_mm,
|
|
"body_center": body_center.tolist(),
|
|
"final_center": transformed_mesh.centroid.tolist(),
|
|
"body_up_axis": int(body_up_axis_idx),
|
|
"ironing": ironing_info,
|
|
}
|
|
|
|
return transformed_mesh, transform, fitting_info
|
|
|
|
|
|
def generate_both_brace_types(
|
|
rigo_type: str,
|
|
output_dir: Path,
|
|
case_id: str,
|
|
cobb_angles: Dict[str, float],
|
|
body_scan_path: Optional[str] = None,
|
|
clearance_mm: float = 8.0
|
|
) -> Dict[str, BraceGenerationResult]:
|
|
"""
|
|
Generate both regular and vase brace types for comparison.
|
|
|
|
Returns:
|
|
Dict with "regular" and "vase" results
|
|
"""
|
|
results = {}
|
|
|
|
# Generate regular brace
|
|
try:
|
|
results["regular"] = generate_glb_brace(
|
|
rigo_type=rigo_type,
|
|
template_type="regular",
|
|
output_dir=output_dir,
|
|
case_id=case_id,
|
|
cobb_angles=cobb_angles,
|
|
body_scan_path=body_scan_path,
|
|
clearance_mm=clearance_mm
|
|
)
|
|
except FileNotFoundError as e:
|
|
results["regular"] = {"error": str(e)}
|
|
|
|
# Generate vase brace
|
|
try:
|
|
results["vase"] = generate_glb_brace(
|
|
rigo_type=rigo_type,
|
|
template_type="vase",
|
|
output_dir=output_dir,
|
|
case_id=case_id,
|
|
cobb_angles=cobb_angles,
|
|
body_scan_path=body_scan_path,
|
|
clearance_mm=clearance_mm
|
|
)
|
|
except FileNotFoundError as e:
|
|
results["vase"] = {"error": str(e)}
|
|
|
|
return results
|
|
|
|
|
|
# Available templates
|
|
AVAILABLE_RIGO_TYPES = ["A1", "A2", "A3", "B1", "B2", "C1", "C2", "E1", "E2"]
|
|
|
|
def list_available_templates() -> Dict[str, list]:
|
|
"""List all available template files."""
|
|
regular = []
|
|
vase = []
|
|
|
|
for rigo_type in AVAILABLE_RIGO_TYPES:
|
|
glb_path, _ = get_template_paths(rigo_type, "regular")
|
|
if glb_path.exists():
|
|
regular.append(rigo_type)
|
|
|
|
glb_path, _ = get_template_paths(rigo_type, "vase")
|
|
if glb_path.exists():
|
|
vase.append(rigo_type)
|
|
|
|
return {
|
|
"regular": regular,
|
|
"vase": vase
|
|
}
|