import torch import io import cv2 import threading import numpy as np import time from pathlib import Path from fastapi import FastAPI, File, UploadFile, Form from fastapi.responses import HTMLResponse from PIL import Image, ImageOps from torchvision import transforms import uvicorn import os from model import BoneAgeModel try: import clip clip_available = True except ImportError: clip_available = False try: import pydicom dicom_available = True except ImportError: dicom_available = False # File size limits (in bytes) MAX_FILE_SIZE = 20 * 1024 * 1024 # 20 MB MIN_FILE_SIZE = 10 * 1024 # 10 KB # Ensure deterministic inference torch.manual_seed(42) np.random.seed(42) torch.use_deterministic_algorithms(True) os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':16:8' # For GPU (if used) app = FastAPI(title="Bone Age Assessment") device = torch.device("cpu") # CPU inference: always use fp32. fp16 on CPU lacks proper kernels and is unstable. model = None for model_file in ["best_model.pth", "best_model_fp16.pth"]: model_path = Path(model_file) if model_path.exists(): try: m = BoneAgeModel() state = torch.load(model_path, map_location="cpu") # If loaded weights are fp16, promote to fp32 for stable CPU inference state = {k: (v.float() if torch.is_tensor(v) and v.dtype == torch.float16 else v) for k, v in state.items()} m.load_state_dict(state) m = m.float() m.eval() model = m print(f"Model loaded: {model_file} (fp32)") break except Exception as e: print(f"Failed to load {model_file}: {e}") continue if model is None: print("No trained model found") # Load CLIP for image validation (hand X-ray detection) clip_model = None clip_preprocess = None if clip_available: try: clip_model, clip_preprocess = clip.load("ViT-B/32", device=device) clip_model.eval() print("CLIP model loaded for image validation") except Exception as e: print(f"Failed to load CLIP: {e}") clip_available = False def _otsu_threshold(gray: np.ndarray) -> int: """Vectorized Otsu's threshold.""" hist, _ = np.histogram(gray, bins=256, range=(0, 256)) hist = hist.astype(float) total = hist.sum() if total == 0: return 127 cumsum = np.cumsum(hist) cum_mean = np.cumsum(hist * np.arange(256)) w_b = cumsum w_f = total - cumsum valid = (w_b > 0) & (w_f > 0) m_b = np.zeros_like(cumsum) m_f = np.zeros_like(cumsum) m_b[valid] = cum_mean[valid] / w_b[valid] m_f[valid] = (cum_mean[-1] - cum_mean[valid]) / w_f[valid] var_between = w_b * w_f * (m_b - m_f) ** 2 var_between[~valid] = 0 return int(np.argmax(var_between)) def auto_crop_hand(image: Image.Image, padding: float = 0.05) -> Image.Image: """Crop to the hand bounding box using Otsu thresholding — removes black borders.""" gray = np.array(image.convert("L")) if gray.std() < 10: return image # near-uniform image, skip thresh = _otsu_threshold(gray) # Normal X-ray: hand brighter than background. Inverted: hand darker. mask = gray < thresh if np.mean(gray) > 128 else gray > thresh rows = np.any(mask, axis=1) cols = np.any(mask, axis=0) if not rows.any() or not cols.any(): return image rmin, rmax = np.where(rows)[0][[0, -1]] cmin, cmax = np.where(cols)[0][[0, -1]] h, w = gray.shape pad_h = int((rmax - rmin) * padding) pad_w = int((cmax - cmin) * padding) rmin = max(0, rmin - pad_h) rmax = min(h, rmax + pad_h + 1) cmin = max(0, cmin - pad_w) cmax = min(w, cmax + pad_w + 1) return image.crop((cmin, rmin, cmax, rmax)) transform = transforms.Compose([ transforms.Resize((512, 512)), # direct resize, no padding (matches training) transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # CLAHE objects are stateful; serialize access for thread safety under FastAPI threadpool _clahe_lock = threading.Lock() _clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) def normalize_xray(image: Image.Image) -> Image.Image: """CLAHE local contrast enhancement.""" gray = np.array(image.convert("L")) # CLAHE: tile-based adaptive histogram equalization (8x8 grid, clip 2.0) with _clahe_lock: gray = _clahe.apply(gray.astype(np.uint8)) return Image.fromarray(gray).convert("RGB") def load_dicom_image(dicom_bytes: bytes) -> Image.Image: """Load DICOM file and convert to PIL Image.""" if not dicom_available: return None try: dicom_data = pydicom.dcmread(io.BytesIO(dicom_bytes)) pixel_array = dicom_data.pixel_array # Handle 16-bit DICOM data if pixel_array.dtype != np.uint8: pixel_array = np.uint8(255 * (pixel_array - pixel_array.min()) / (pixel_array.max() - pixel_array.min())) # Handle multi-frame (take first frame) if len(pixel_array.shape) == 3: pixel_array = pixel_array[0] return Image.fromarray(pixel_array).convert("L") except Exception as e: print(f"DICOM loading error: {e}") return None def validate_hand_xray(image: Image.Image, log_timing: bool = False, debug: bool = True) -> tuple[bool, str, float, dict]: """Validate that image is a hand X-ray using CLIP. Returns (is_valid, message, confidence, timing_dict).""" timings = {} if not clip_available or clip_model is None: return True, "CLIP validation skipped", 0.0, timings try: # Preprocess image for CLIP t0 = time.perf_counter() image_input = clip_preprocess(image).unsqueeze(0).to(device) timings['clip_preprocess'] = (time.perf_counter() - t0) * 1000 # Descriptions to match against descriptions = [ "a hand X-ray image", "a pediatric hand X-ray", "a chest X-ray", "a leg X-ray", "a spine X-ray", "a CT scan", "an MRI image", "a photograph", "a random image" ] t0 = time.perf_counter() with torch.no_grad(): # Encode image image_features = clip_model.encode_image(image_input) image_features /= image_features.norm(dim=-1, keepdim=True) # Encode text descriptions text_tokens = clip.tokenize(descriptions).to(device) text_features = clip_model.encode_text(text_tokens) text_features /= text_features.norm(dim=-1, keepdim=True) # Calculate similarity scores similarity = (image_features @ text_features.T).squeeze(0) scores = similarity.softmax(dim=0).cpu().numpy() timings['clip_inference'] = (time.perf_counter() - t0) * 1000 # Hand X-ray descriptions (indices 0, 1) hand_xray_score = float(scores[0] + scores[1]) # Non-hand/non-xray descriptions (indices 2-8) non_hand_score = float(scores[2:].sum()) # Confidence: how confident we are it's a hand X-ray confidence = hand_xray_score - non_hand_score if log_timing or debug: total = sum(timings.values()) print(f"[CLIP TIMING] Preprocess: {timings['clip_preprocess']:.1f}ms | Inference: {timings['clip_inference']:.1f}ms | Total: {total:.1f}ms") if debug: print("[CLIP DEBUG] Individual scores:") for i, desc in enumerate(descriptions): print(f" [{i}] {desc}: {scores[i]:.4f}") print(f"[CLIP DEBUG] Hand X-ray score (0+1): {hand_xray_score:.4f} ({hand_xray_score:.1%})") print(f"[CLIP DEBUG] Non-hand score (2-8): {non_hand_score:.4f}") print(f"[CLIP DEBUG] Confidence (hand - non-hand): {confidence:.4f}") print(f"[CLIP DEBUG] Decision: hand_xray > non_hand? {hand_xray_score > non_hand_score}") # Threshold: hand X-ray score must be higher than non-hand categories if hand_xray_score > non_hand_score: return True, f"Valid hand X-ray detected (confidence: {hand_xray_score:.1%})", hand_xray_score, timings else: return False, "Image does not appear to be a hand X-ray. Upload a PA hand radiograph for interpretation.", hand_xray_score, timings except Exception as e: print(f"CLIP validation error: {e}") return True, "Validation check failed, proceeding", 0.0, timings def predict_bone_age(image: Image.Image, is_male: bool, log_timing: bool = False) -> tuple[float, dict]: """Inference with test-time augmentation (original + flip) for better accuracy. Returns: tuple: (predicted_age_in_months, timing_dict) """ timings = {} # EXIF transpose t0 = time.perf_counter() image = ImageOps.exif_transpose(image) timings['exif_transpose'] = (time.perf_counter() - t0) * 1000 # Auto-crop hand ROI t0 = time.perf_counter() image = auto_crop_hand(image) timings['auto_crop'] = (time.perf_counter() - t0) * 1000 # CLAHE normalization t0 = time.perf_counter() image = normalize_xray(image) timings['clahe_normalize'] = (time.perf_counter() - t0) * 1000 gender_tensor = torch.tensor([1.0 if is_male else 0.0]).to(device) # TTA: average predictions from original and horizontally flipped image t0 = time.perf_counter() preds = [] for img in (image, image.transpose(Image.FLIP_LEFT_RIGHT)): img_tensor = transform(img).unsqueeze(0).to(device) with torch.no_grad(): preds.append(model(img_tensor, gender_tensor).item()) timings['model_inference_tta'] = (time.perf_counter() - t0) * 1000 result = sum(preds) / len(preds) if log_timing: total = sum(timings.values()) print(f"[TIMING] EXIF: {timings['exif_transpose']:.1f}ms | Auto-crop: {timings['auto_crop']:.1f}ms | " f"CLAHE: {timings['clahe_normalize']:.1f}ms | Model: {timings['model_inference_tta']:.1f}ms | " f"Total: {total:.1f}ms") return result, timings @app.get("/", response_class=HTMLResponse) async def home(): return """
Upload a hand X-ray to estimate skeletal maturity