braceiqmed/brace-generator/data_models.py

"""
Data models for unified spine landmark representation.
This is the "glue" that connects different model outputs to the brace generator.
"""
from dataclasses import dataclass, field
from typing import Optional, List, Dict, Any
import numpy as np


@dataclass
class VertebraLandmark:
    """
    Unified representation of a single vertebra's landmarks.
    All coordinates are in pixels (can be converted to mm with pixel_spacing).
    """
    # Vertebra level identifier (e.g., "T1", "T4", "L1", etc.) - None if unknown
    level: Optional[str] = None
    
    # Center point of vertebra [x, y] in pixels
    centroid_px: np.ndarray = field(default_factory=lambda: np.zeros(2))
    
    # Four corner points [top_left, top_right, bottom_right, bottom_left] shape (4, 2)
    corners_px: Optional[np.ndarray] = None
    
    # Upper endplate points [left, right] shape (2, 2)
    endplate_upper_px: Optional[np.ndarray] = None
    
    # Lower endplate points [left, right] shape (2, 2)  
    endplate_lower_px: Optional[np.ndarray] = None
    
    # Orientation angle of vertebra in degrees (tilt in coronal plane)
    orientation_deg: Optional[float] = None
    
    # Detection confidence (0-1)
    confidence: float = 1.0
    
    # Additional metadata from source model
    meta: Optional[Dict[str, Any]] = None
    
    def compute_orientation(self) -> float:
        """Compute vertebra orientation from corners or endplates."""
        if self.orientation_deg is not None:
            return self.orientation_deg
            
        # Try to compute from upper endplate
        if self.endplate_upper_px is not None:
            left, right = self.endplate_upper_px[0], self.endplate_upper_px[1]
            dx = right[0] - left[0]
            dy = right[1] - left[1]
            angle = np.degrees(np.arctan2(dy, dx))
            self.orientation_deg = angle
            return angle
            
        # Try to compute from corners (top-left to top-right)
        if self.corners_px is not None:
            top_left, top_right = self.corners_px[0], self.corners_px[1]
            dx = top_right[0] - top_left[0]
            dy = top_right[1] - top_left[1]
            angle = np.degrees(np.arctan2(dy, dx))
            self.orientation_deg = angle
            return angle
            
        return 0.0
    
    def compute_centroid(self) -> np.ndarray:
        """Compute centroid from corners if not set."""
        if self.corners_px is not None and np.all(self.centroid_px == 0):
            self.centroid_px = np.mean(self.corners_px, axis=0)
        return self.centroid_px


@dataclass
class Spine2D:
    """
    Complete 2D spine representation from an X-ray.
    Contains all detected vertebrae and computed angles.
    """
    # List of vertebrae, ordered from top (C7/T1) to bottom (L5/S1)
    vertebrae: List[VertebraLandmark] = field(default_factory=list)
    
    # Pixel spacing in mm [sx, sy] - from DICOM if available
    pixel_spacing_mm: Optional[np.ndarray] = None
    
    # Original image shape (height, width)
    image_shape: Optional[tuple] = None
    
    # Computed Cobb angles in degrees (individual fields)
    cobb_pt: Optional[float] = None  # Proximal Thoracic
    cobb_mt: Optional[float] = None  # Main Thoracic  
    cobb_tl: Optional[float] = None  # Thoracolumbar/Lumbar
    
    # Cobb angles as dictionary (alternative format)
    cobb_angles: Optional[Dict[str, float]] = None  # {'PT': angle, 'MT': angle, 'TL': angle}
    
    # Curve type: "S" (double curve) or "C" (single curve) or "Normal"
    curve_type: Optional[str] = None
    
    # Rigo-Chêneau classification
    rigo_type: Optional[str] = None  # A1, A2, A3, B1, B2, C1, C2, E1, E2, Normal
    rigo_description: Optional[str] = None  # Detailed description
    
    # Source model that generated this data
    source_model: Optional[str] = None
    
    # Additional metadata
    meta: Optional[Dict[str, Any]] = None
    
    def get_cobb_angles(self) -> Dict[str, float]:
        """Get Cobb angles as dictionary, preferring computed individual fields."""
        # Prefer individual fields (set by compute_cobb_angles) over dictionary
        # This ensures consistency between displayed values and classification
        if self.cobb_pt is not None or self.cobb_mt is not None or self.cobb_tl is not None:
            return {
                'PT': self.cobb_pt or 0.0,
                'MT': self.cobb_mt or 0.0,
                'TL': self.cobb_tl or 0.0
            }
        if self.cobb_angles is not None:
            return self.cobb_angles
        return {'PT': 0.0, 'MT': 0.0, 'TL': 0.0}
    
    def get_centroids(self) -> np.ndarray:
        """Get array of all vertebra centroids, shape (N, 2)."""
        centroids = []
        for v in self.vertebrae:
            v.compute_centroid()
            centroids.append(v.centroid_px)
        return np.array(centroids, dtype=np.float32)
    
    def get_orientations(self) -> np.ndarray:
        """Get array of all vertebra orientations in degrees, shape (N,)."""
        return np.array([v.compute_orientation() for v in self.vertebrae], dtype=np.float32)
    
    def to_mm(self, coords_px: np.ndarray) -> np.ndarray:
        """Convert pixel coordinates to millimeters."""
        if self.pixel_spacing_mm is None:
            # Default assumption: 0.25 mm/pixel (typical for spine X-rays)
            spacing = np.array([0.25, 0.25])
        else:
            spacing = self.pixel_spacing_mm
        return coords_px * spacing
    
    def sort_vertebrae(self):
        """Sort vertebrae by vertical position (top to bottom)."""
        self.vertebrae.sort(key=lambda v: float(v.centroid_px[1]))


@dataclass
class BraceConfig:
    """
    Configuration parameters for brace generation.
    """
    # Brace dimensions
    brace_height_mm: float = 400.0  # Total height of brace
    wall_thickness_mm: float = 4.0  # Shell thickness
    
    # Torso shape parameters (for average body mode)
    torso_width_mm: float = 280.0   # Left-right diameter at widest
    torso_depth_mm: float = 200.0   # Front-back diameter at widest
    
    # Correction parameters
    pressure_strength_mm: float = 15.0  # Max indentation at apex
    pressure_spread_deg: float = 45.0   # Angular spread of pressure zone
    
    # Mesh resolution
    n_vertical_slices: int = 100    # Number of cross-sections
    n_circumference_points: int = 72  # Points per cross-section (every 5°)
    
    # Opening (for brace accessibility)
    front_opening_deg: float = 60.0  # Angular width of front opening (0 = closed)
    
    # Mode
    use_body_scan: bool = False     # True = use 3D body scan, False = average body
    body_scan_path: Optional[str] = None  # Path to body scan mesh
    
    # Scale
    pixel_spacing_mm: Optional[np.ndarray] = None  # Override pixel spacing