Initial commit - BraceIQMed platform with frontend, API, and brace generator
This commit is contained in:
906
brace-generator/glb_generator.py
Normal file
906
brace-generator/glb_generator.py
Normal file
@@ -0,0 +1,906 @@
|
||||
"""
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user