Files
braceiqmed/brace-generator/schemas.py

126 lines
3.6 KiB
Python

"""
Pydantic schemas for API request/response validation.
"""
from typing import Optional, List, Dict, Any
from pydantic import BaseModel, Field
from enum import Enum
class ExperimentType(str, Enum):
"""Available brace generation experiments."""
STANDARD = "standard" # Original pipeline
EXPERIMENT_3 = "experiment_3" # Research-based adaptive
class BraceConfigRequest(BaseModel):
"""Brace configuration parameters."""
brace_height_mm: float = Field(default=400.0, ge=200, le=600)
torso_width_mm: float = Field(default=280.0, ge=150, le=400)
torso_depth_mm: float = Field(default=200.0, ge=100, le=350)
wall_thickness_mm: float = Field(default=4.0, ge=2, le=10)
pressure_strength_mm: float = Field(default=15.0, ge=0, le=30)
class AnalyzeRequest(BaseModel):
"""Request to analyze X-ray and generate brace."""
s3_key: Optional[str] = Field(None, description="S3 key of uploaded X-ray image")
case_id: Optional[str] = Field(None, description="Case ID for organizing outputs")
experiment: ExperimentType = Field(default=ExperimentType.EXPERIMENT_3)
config: Optional[BraceConfigRequest] = None
# Output options
save_visualization: bool = Field(default=True)
save_landmarks: bool = Field(default=True)
output_format: str = Field(default="stl", description="stl, ply, or both")
class AnalyzeFromUrlRequest(BaseModel):
"""Request with direct image URL."""
image_url: str = Field(..., description="URL to download X-ray image from")
case_id: Optional[str] = Field(None)
experiment: ExperimentType = Field(default=ExperimentType.EXPERIMENT_3)
config: Optional[BraceConfigRequest] = None
save_visualization: bool = True
save_landmarks: bool = True
output_format: str = "stl"
class Vertebra(BaseModel):
"""Single vertebra data."""
level: str
centroid_px: List[float]
orientation_deg: Optional[float] = None
confidence: Optional[float] = None
corners_px: Optional[List[List[float]]] = None
class CobbAngles(BaseModel):
"""Cobb angle measurements."""
PT: float = Field(..., description="Proximal Thoracic angle")
MT: float = Field(..., description="Main Thoracic angle")
TL: float = Field(..., description="Thoracolumbar angle")
class RigoClassification(BaseModel):
"""Rigo-Chêneau classification result."""
type: str
description: str
curve_pattern: Optional[str] = None
class DeformationReport(BaseModel):
"""Patch-based deformation report (Experiment 3)."""
patch_grid: str
deformations: Optional[List[List[float]]] = None
zones: Optional[List[Dict[str, Any]]] = None
class AnalysisResult(BaseModel):
"""Complete analysis result."""
case_id: Optional[str] = None
experiment: str
# Input
input_image: str
# Detection results
model_used: str
vertebrae_detected: int
vertebrae: Optional[List[Vertebra]] = None
# Measurements
cobb_angles: CobbAngles
curve_type: str
# Classification
rigo_classification: RigoClassification
# Brace mesh info
mesh_vertices: int
mesh_faces: int
# Deformation (Experiment 3)
deformation_report: Optional[DeformationReport] = None
# Output URLs/paths
outputs: Dict[str, str] = Field(default_factory=dict)
# Timing
processing_time_ms: float
class HealthResponse(BaseModel):
"""Health check response."""
status: str
device: str
cuda_available: bool
model_loaded: bool
gpu_name: Optional[str] = None
gpu_memory_mb: Optional[int] = None
class ErrorResponse(BaseModel):
"""Error response."""
error: str
detail: Optional[str] = None