Update app.py

app.py

@@ -6,9 +6,10 @@ 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.layers import Input, GlobalAveragePooling2D, Dense
+from tensorflow.keras.applications import DenseNet121
 from tensorflow.keras.preprocessing.image import img_to_array
-from tensorflow.keras.applications.densenet import DenseNet121, preprocess_input
+from tensorflow.keras.applications.densenet import preprocess_input
 import numpy as np
 from PIL import Image
 import matplotlib.pyplot as plt
@@ -27,7 +28,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.3.0"
+    version="3.0.0"
 )
 
 # Rate limiter setup
@@ -57,37 +58,47 @@ 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))
+def build_custom_model():
+    """Build model with correct output shape matching your weights"""
+    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)
+    x = GlobalAveragePooling2D()(x)
+    # Match the output shape in your pretrained weights (14 classes)
+    predictions = Dense(14, activation='sigmoid')(x)
     return Model(inputs=base_model.input, outputs=predictions)
 
-def load_model_with_fallback():
-    """Attempt multiple strategies to load the model"""
+def load_model_with_retry():
+    """Enhanced model loading with shape compatibility handling"""
     try:
-        # Strategy 1: Try direct loading
-        model = load_model('Densenet.h5', compile=False)
+        # First try loading with custom architecture
+        model = build_custom_model()
         model.load_weights('pretrained_model.h5')
         return model
     except Exception as e:
-        print(f"Direct loading failed: {e}")
-        
+        print(f"Loading with custom architecture failed: {e}")
         try:
-            # Strategy 2: Build architecture and load weights
-            model = build_model()
-            model.load_weights('pretrained_model.h5')
+            # Fallback to direct loading with compile=False
+            model = load_model('Densenet.h5', compile=False)
+            # Ensure output layer matches our class names
+            if model.layers[-1].output_shape[-1] != len(class_names):
+                print("Adjusting output layer to match class names")
+                x = model.layers[-2].output
+                predictions = Dense(len(class_names), activation='sigmoid')(x)
+                model = Model(inputs=model.input, outputs=predictions)
             return model
         except Exception as e:
-            print(f"Architecture rebuild failed: {e}")
-            raise RuntimeError("All model loading strategies failed")
+            print(f"All loading attempts failed: {e}")
+            raise RuntimeError(f"Could not load model: {str(e)}")
 
 # Load model
 try:
-    model = load_model_with_fallback()
+    model = load_model_with_retry()
     print("✅ Model loaded successfully!")
+    print(f"Model output shape: {model.output_shape}")
 except Exception as e:
     print(f"❌ Model loading failed: {e}")
     raise
@@ -133,10 +144,11 @@ def generate_gradcam(img):
     return Image.blend(original_img, heatmap_img, 0.5)
 
 def process_predictions(predictions):
-    """Format predictions with top 4 classes"""
+    """Format predictions with top classes"""
     decoded = []
     for pred in predictions:
-        top_indices = pred.argsort()[-4:][::-1]
+        # Get indices sorted by probability (descending)
+        top_indices = np.argsort(pred)[::-1][:len(class_names)]
         decoded.append([(class_names[i], float(pred[i])) for i in top_indices])
     return decoded
 
@@ -162,7 +174,8 @@ async def health_check():
     return {
         "status": "healthy" if model else "unhealthy",
         "timestamp": datetime.now().isoformat(),
-        "model_loaded": bool(model)
+        "model_loaded": bool(model),
+        "model_output_shape": str(model.output_shape) if model else "N/A"
     }
 
 @app.get("/model/classes")
@@ -172,6 +185,7 @@ 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"""
     if not file.content_type.startswith('image/'):
         raise HTTPException(400, "Only image files are accepted")
     
@@ -215,6 +229,7 @@ async def analyze_image(request: Request, file: UploadFile = File(...)):
 
 @app.get("/static/heatmap/{session_id}")
 async def get_heatmap(session_id: str):
+    """Retrieve Grad-CAM visualization"""
     if session_id not in heatmap_store:
         raise HTTPException(404, "Session expired or invalid")
     return StreamingResponse(
@@ -225,10 +240,12 @@ 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",
         "classes": len(class_names),
+        "output_shape": str(model.output_shape),
         "gradcam_layer": layer_name,
         "rate_limit": "5 requests/minute"
     }