Update app.py

app.py

@@ -6,8 +6,9 @@ from slowapi import Limiter
 from slowapi.util import get_remote_address
 import tensorflow as tf
 from tensorflow.keras.models import Model, load_model
+from tensorflow.keras.layers import Input
 from tensorflow.keras.preprocessing.image import img_to_array
-from tensorflow.keras.applications.densenet import preprocess_input
+from tensorflow.keras.applications.densenet import DenseNet121, preprocess_input
 import numpy as np
 from PIL import Image
 import matplotlib.pyplot as plt
@@ -26,7 +27,7 @@ HEATMAP_EXPIRY = 300  # 5 minutes in seconds
 app = FastAPI(
     title="ChexNet Medical Imaging API",
     description="API for chest X-ray analysis with Grad-CAM visualization",
-    version="2.1.0"
+    version="2.3.0"
 )
 
 # Rate limiter setup
@@ -48,13 +49,6 @@ app.add_middleware(
 # Session storage for heatmaps
 heatmap_store: Dict[str, dict] = {}
 
-# Load model
-try:
-    model = load_model('Densenet.h5')
-    model.load_weights("pretrained_model.h5")
-except Exception as e:
-    raise RuntimeError(f"Failed to load model: {str(e)}")
-
 # Model configuration
 layer_name = 'conv5_block16_concat'
 class_names = [
@@ -63,6 +57,41 @@ class_names = [
     'Pneumonia', 'Fibrosis', 'Edema', 'Consolidation', 'No Finding'
 ]
 
+def build_model():
+    """Build DenseNet121 model architecture"""
+    base_model = DenseNet121(weights=None, include_top=False, input_shape=(None, None, 3))
+    x = base_model.output
+    x = tf.keras.layers.GlobalAveragePooling2D()(x)
+    predictions = tf.keras.layers.Dense(len(class_names), activation='sigmoid')(x)
+    return Model(inputs=base_model.input, outputs=predictions)
+
+def load_model_with_fallback():
+    """Attempt multiple strategies to load the model"""
+    try:
+        # Strategy 1: Try direct loading
+        model = load_model('Densenet.h5', compile=False)
+        model.load_weights('pretrained_model.h5')
+        return model
+    except Exception as e:
+        print(f"Direct loading failed: {e}")
+        
+        try:
+            # Strategy 2: Build architecture and load weights
+            model = build_model()
+            model.load_weights('pretrained_model.h5')
+            return model
+        except Exception as e:
+            print(f"Architecture rebuild failed: {e}")
+            raise RuntimeError("All model loading strategies failed")
+
+# Load model
+try:
+    model = load_model_with_fallback()
+    print("✅ Model loaded successfully!")
+except Exception as e:
+    print(f"❌ Model loading failed: {e}")
+    raise
+
 def cleanup_expired_heatmaps():
     """Remove heatmaps older than HEATMAP_EXPIRY seconds"""
     now = datetime.now()
@@ -73,7 +102,7 @@ def cleanup_expired_heatmaps():
     for sid in expired:
         del heatmap_store[sid]
 
-def generate_gradcam(model, img, layer_name):
+def generate_gradcam(img):
     """Generate Grad-CAM heatmap overlay"""
     img_array = img_to_array(img)
     img_array = np.expand_dims(img_array, axis=0)
@@ -103,12 +132,12 @@ def generate_gradcam(model, img, layer_name):
     heatmap_img = heatmap_img.resize(original_img.size)
     return Image.blend(original_img, heatmap_img, 0.5)
 
-def process_predictions(predictions, class_labels):
+def process_predictions(predictions):
     """Format predictions with top 4 classes"""
     decoded = []
     for pred in predictions:
-        top_indices = pred.argsort()[-4:][::-1]  # Top 4 predictions
-        decoded.append([(class_labels[i], float(pred[i])) for i in top_indices])
+        top_indices = pred.argsort()[-4:][::-1]
+        decoded.append([(class_names[i], float(pred[i])) for i in top_indices])
     return decoded
 
 def preprocess_image(file_bytes):
@@ -119,14 +148,21 @@ def preprocess_image(file_bytes):
 
 @app.get("/", include_in_schema=False)
 async def root():
-    return {"message": "ChexNet API is operational", "docs": "/docs"}
+    return {
+        "message": "ChexNet API is operational",
+        "endpoints": {
+            "docs": "/docs",
+            "health": "/health",
+            "analyze": "POST /analyze"
+        }
+    }
 
 @app.get("/health")
 async def health_check():
     return {
-        "status": "healthy",
-        "model_loaded": model is not None,
-        "timestamp": datetime.now().isoformat()
+        "status": "healthy" if model else "unhealthy",
+        "timestamp": datetime.now().isoformat(),
+        "model_loaded": bool(model)
     }
 
 @app.get("/model/classes")
@@ -136,26 +172,15 @@ async def get_class_names():
 @app.post("/analyze")
 @limiter.limit("5/minute")
 async def analyze_image(request: Request, file: UploadFile = File(...)):
-    """
-    Analyze chest X-ray image and return predictions with Grad-CAM visualization
-    
-    Parameters:
-    - file: Upload JPEG/PNG image (max 10MB)
-    
-    Returns:
-    - predictions: Top 4 diagnoses with confidence scores
-    - heatmap_url: URL to retrieve Grad-CAM visualization
-    """
-    # Validate input
     if not file.content_type.startswith('image/'):
         raise HTTPException(400, "Only image files are accepted")
     
     if file.size > MAX_FILE_SIZE:
-        raise HTTPException(413, f"Maximum file size is {MAX_FILE_SIZE//(1024*1024)}MB")
+        raise HTTPException(413, f"Max file size is {MAX_FILE_SIZE//(1024*1024)}MB")
     
     try:
-        # Process image
-        img = preprocess_image(await file.read())
+        contents = await file.read()
+        img = preprocess_image(contents)
         
         # Prepare input tensor
         img_array = img_to_array(img)
@@ -164,10 +189,10 @@ async def analyze_image(request: Request, file: UploadFile = File(...)):
         
         # Get predictions
         predictions = model.predict(img_array)
-        decoded = process_predictions(predictions, class_names)
+        decoded = process_predictions(predictions)
         
         # Generate Grad-CAM
-        heatmap = generate_gradcam(model, img, layer_name)
+        heatmap = generate_gradcam(img)
         
         # Store heatmap with session ID
         session_id = str(uuid.uuid4())
@@ -185,16 +210,13 @@ async def analyze_image(request: Request, file: UploadFile = File(...)):
             "predictions": decoded[0],
             "heatmap_url": f"{request.base_url}static/heatmap/{session_id}"
         }
-        
     except Exception as e:
         raise HTTPException(500, f"Processing failed: {str(e)}")
 
 @app.get("/static/heatmap/{session_id}")
 async def get_heatmap(session_id: str):
-    """Retrieve Grad-CAM visualization by session ID"""
     if session_id not in heatmap_store:
         raise HTTPException(404, "Session expired or invalid")
-    
     return StreamingResponse(
         io.BytesIO(heatmap_store[session_id]['image']),
         media_type="image/png",
@@ -203,7 +225,6 @@ async def get_heatmap(session_id: str):
 
 @app.get("/model/info")
 async def model_info():
-    """Get model metadata"""
     return {
         "model_type": "DenseNet121",
         "input_size": "540x540",
@@ -212,16 +233,15 @@ async def model_info():
         "rate_limit": "5 requests/minute"
     }
 
-# Error handlers
 @app.exception_handler(HTTPException)
-async def handle_http_exception(request, exc):
+async def http_handler(request: Request, exc: HTTPException):
     return JSONResponse(
         status_code=exc.status_code,
         content={"error": exc.detail}
     )
 
 @app.exception_handler(Exception)
-async def handle_generic_exception(request, exc):
+async def generic_handler(request: Request, exc: Exception):
     return JSONResponse(
         status_code=500,
         content={"error": "Internal server error"}