File size: 6,461 Bytes
14622e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | import os
import io
import cv2
import logging
import numpy as np
# ==========================================
# 0. PRODUCTION LOGGING & ENVIRONMENT SETUP
# ==========================================
# Set TensorFlow variables BEFORE importing TF to suppress warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
import tensorflow as tf
from tensorflow import keras
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
# Configure production logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("AiroDx-API")
# ==========================================
# 1. CONFIGURATION & PATHS
# ==========================================
# Use Environment Variables for paths (Allows Docker to change them without changing code)
MODEL_PATH = os.getenv("MODEL_PATH", "./model/remedis_multiview_final.keras")
REMEDIS_BASE_PATH = os.getenv("REMEDIS_BASE_PATH", "./model/cxr-50x1-remedis-m")
IMAGE_SIZE = (448, 448)
TARGET_DISEASES = [
'Atelectasis', 'Cardiomegaly', 'Consolidation', 'Emphysema', 'Hernia',
'Infiltrates', 'Mass', 'Nodule', 'Pleural Effusion', 'Pleural Thickening',
'Pneumonia', 'Pneumothorax', 'Pulmonary Edema', 'Pulmonary Fibrosis'
]
# ==========================================
# 2. REDEFINE CUSTOM LAYER
# ==========================================
@tf.keras.utils.register_keras_serializable()
class REMEDISBackbone(keras.layers.Layer):
def __init__(self, model_path=None, **kwargs):
# Accept legacy serialized model_path, but force the local deployment path.
super().__init__(**kwargs)
self.model_path = REMEDIS_BASE_PATH
self._remedis_model = None
def build(self, input_shape):
if self._remedis_model is None:
if not os.path.exists(self.model_path):
logger.error(f"Base REMEDIS model not found at {self.model_path}")
raise FileNotFoundError(f"Base REMEDIS model not found at {self.model_path}")
self._remedis_model = tf.saved_model.load(self.model_path)
super().build(input_shape)
def call(self, inputs, training=None):
return self._remedis_model(inputs)
def compute_output_shape(self, input_shape):
return (input_shape[0], 14, 14, 2048)
def get_config(self):
config = super().get_config()
return config
@classmethod
def from_config(cls, config):
# Remove stale serialized path (e.g., Kaggle) before layer construction.
config.pop("model_path", None)
return cls(**config)
# ==========================================
# 3. INITIALIZE APP & LOAD MODEL
# ==========================================
app = FastAPI(
title="AiroDx Multi-View Diagnostic API",
description="Multimodal diagnostic support using Siamese REMEDIS network.",
version="1.0.0"
)
# CORS setup for production (Can be restricted via ENV variable)
ALLOWED_ORIGINS = os.getenv("ALLOWED_ORIGINS", "*").split(",")
app.add_middleware(
CORSMiddleware,
allow_origins=ALLOWED_ORIGINS,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
logger.info(f"Loading weights from {MODEL_PATH}...")
try:
model = keras.models.load_model(
MODEL_PATH,
custom_objects={"REMEDISBackbone": REMEDISBackbone}
)
logger.info("Model loaded successfully into memory!")
except Exception as e:
logger.critical(f"Failed to load model: {e}")
model = None
# ==========================================
# 4. HEALTH CHECK ENDPOINT (For Cloud Hosting)
# ==========================================
@app.get("/health")
async def health_check():
"""Endpoint used by load balancers and Docker to ensure API is alive."""
if model is None:
raise HTTPException(status_code=503, detail="Model failed to load.")
return {"status": "healthy", "model_ready": True}
# ==========================================
# 5. PREPROCESSING FUNCTION
# ==========================================
async def process_image(file: UploadFile) -> np.ndarray:
try:
# Read the file bytes
contents = await file.read()
nparr = np.frombuffer(contents, np.uint8)
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
if img is None:
raise ValueError(f"Invalid image file: {file.filename}")
# Exact same preprocessing used in training
img = cv2.resize(img, IMAGE_SIZE)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = img.astype('float32') / 255.0
# Add batch dimension: (448, 448, 3) -> (1, 448, 448, 3)
return np.expand_dims(img, axis=0)
except Exception as e:
logger.error(f"Image processing error: {str(e)}")
raise HTTPException(status_code=400, detail=f"Error processing image: {str(e)}")
# ==========================================
# 6. PREDICTION ENDPOINT
# ==========================================
@app.post("/predict")
async def predict(
frontal_image: UploadFile = File(..., description="The PA/AP Frontal X-ray"),
lateral_image: UploadFile = File(..., description="The Lateral X-ray")
):
logger.info(f"Received prediction request for files: {frontal_image.filename}, {lateral_image.filename}")
# 1. Process both images
img_pa = await process_image(frontal_image)
img_l = await process_image(lateral_image)
# 2. Run inference
predictions = model.predict({
"pa_image": img_pa,
"l_image": img_l
})
# 3. Format results
probabilities = predictions[0].tolist()
results = []
for disease, prob in zip(TARGET_DISEASES, probabilities):
results.append({
"disease": disease,
"probability": round(prob, 4),
"flagged": bool(prob > 0.5) # You can adjust this threshold later based on your F1 tuning
})
# Sort results by probability (highest first)
results.sort(key=lambda x: x["probability"], reverse=True)
logger.info("Prediction processed successfully.")
return {
"status": "success",
"predictions": results
} |