Update app.py

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException, Request
+from fastapi import FastAPI, UploadFile, File, HTTPException, Request, BackgroundTasks
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
@@ -6,7 +6,7 @@ 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, GlobalAveragePooling2D, Dense
+from tensorflow.keras.layers import GlobalAveragePooling2D, Dense
 from tensorflow.keras.applications import DenseNet121
 from tensorflow.keras.preprocessing.image import img_to_array
 from tensorflow.keras.applications.densenet import preprocess_input
@@ -18,23 +18,29 @@ import io
 import uuid
 from typing import Dict
 from datetime import datetime, timedelta
+from pathlib import Path
 import os
 
 # Configuration
 MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB
 HEATMAP_EXPIRY = 300  # 5 minutes in seconds
+PORT = 7860  # Hugging Face Spaces requires port 7860
 
 # Initialize FastAPI with rate limiting
 app = FastAPI(
     title="ChexNet Medical Imaging API",
     description="API for chest X-ray analysis with Grad-CAM visualization",
-    version="3.0.0"
+    version="4.0.0"
 )
 
 # Rate limiter setup
 limiter = Limiter(key_func=get_remote_address)
 app.state.limiter = limiter
 
+# Create static/heatmap directory
+heatmap_dir = Path("static/heatmap")
+heatmap_dir.mkdir(parents=True, exist_ok=True)
+
 # Mount static files
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
@@ -47,7 +53,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Session storage for heatmaps
+# Session storage for heatmaps (now using file system)
 heatmap_store: Dict[str, dict] = {}
 
 # Model configuration
@@ -58,8 +64,8 @@ class_names = [
     'Pneumonia', 'Fibrosis', 'Edema', 'Consolidation', 'No Finding'
 ]
 
-def build_custom_model():
-    """Build model with correct output shape matching your weights"""
+def build_model():
+    """Build DenseNet121 model with correct output shape"""
     base_model = DenseNet121(
         weights=None,
         include_top=False,
@@ -67,54 +73,50 @@ def build_custom_model():
     )
     x = base_model.output
     x = GlobalAveragePooling2D()(x)
-    # Match the output shape in your pretrained weights (14 classes)
-    predictions = Dense(14, activation='sigmoid')(x)
+    predictions = Dense(14, activation='sigmoid')(x)  # 14 classes in pretrained weights
     return Model(inputs=base_model.input, outputs=predictions)
 
-def load_model_with_retry():
-    """Enhanced model loading with shape compatibility handling"""
+def load_model_with_fallback():
+    """Robust model loading with multiple fallback strategies"""
     try:
-        # First try loading with custom architecture
-        model = build_custom_model()
+        # Strategy 1: Build model and load weights
+        model = build_model()
         model.load_weights('pretrained_model.h5')
         return model
     except Exception as e:
-        print(f"Loading with custom architecture failed: {e}")
+        print(f"Primary loading failed: {e}")
         try:
-            # Fallback to direct loading with compile=False
+            # Strategy 2: Try direct loading
             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"All loading attempts failed: {e}")
-            raise RuntimeError(f"Could not load model: {str(e)}")
+            print(f"Fallback loading failed: {e}")
+            raise RuntimeError("All model loading strategies failed")
 
 # Load model
 try:
-    model = load_model_with_retry()
+    model = load_model_with_fallback()
     print("✅ Model loaded successfully!")
+    print(f"Model input shape: {model.input_shape}")
     print(f"Model output shape: {model.output_shape}")
 except Exception as e:
     print(f"❌ Model loading failed: {e}")
     raise
 
-def cleanup_expired_heatmaps():
-    """Remove heatmaps older than HEATMAP_EXPIRY seconds"""
+async def cleanup_old_heatmaps():
+    """Delete heatmap files older than HEATMAP_EXPIRY seconds"""
     now = datetime.now()
-    expired = [
-        sid for sid, data in heatmap_store.items()
-        if (now - data['timestamp']).total_seconds() > HEATMAP_EXPIRY
-    ]
-    for sid in expired:
-        del heatmap_store[sid]
+    for file in heatmap_dir.glob("*.png"):
+        file_time = datetime.fromtimestamp(file.stat().st_mtime)
+        if (now - file_time) > timedelta(seconds=HEATMAP_EXPIRY):
+            try:
+                file.unlink()
+                print(f"Deleted expired heatmap: {file.name}")
+            except Exception as e:
+                print(f"Error deleting {file.name}: {e}")
 
 def generate_gradcam(img):
-    """Generate Grad-CAM heatmap overlay"""
+    """Generate Grad-CAM heatmap visualization"""
     img_array = img_to_array(img)
     img_array = np.expand_dims(img_array, axis=0)
     img_array = preprocess_input(img_array)
@@ -144,24 +146,23 @@ def generate_gradcam(img):
     return Image.blend(original_img, heatmap_img, 0.5)
 
 def process_predictions(predictions):
-    """Format predictions with top classes"""
+    """Format model predictions with confidence scores"""
     decoded = []
     for pred in predictions:
-        # 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
 
 def preprocess_image(file_bytes):
-    """Convert uploaded file to processed image array"""
+    """Process uploaded image file"""
     img = cv2.imdecode(np.frombuffer(file_bytes, np.uint8), cv2.IMREAD_COLOR)
     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     return cv2.resize(img, (540, 540), interpolation=cv2.INTER_AREA)
 
-@app.get("/", include_in_schema=False)
+@app.get("/")
 async def root():
     return {
-        "message": "ChexNet API is operational",
+        "message": "ChexNet API is running",
         "endpoints": {
             "docs": "/docs",
             "health": "/health",
@@ -172,10 +173,11 @@ async def root():
 @app.get("/health")
 async def health_check():
     return {
-        "status": "healthy" if model else "unhealthy",
+        "status": "healthy",
+        "model_loaded": True,
         "timestamp": datetime.now().isoformat(),
-        "model_loaded": bool(model),
-        "model_output_shape": str(model.output_shape) if model else "N/A"
+        "port": PORT,
+        "heatmap_files": len(list(heatmap_dir.glob("*.png")))
     }
 
 @app.get("/model/classes")
@@ -184,13 +186,17 @@ async def get_class_names():
 
 @app.post("/analyze")
 @limiter.limit("5/minute")
-async def analyze_image(request: Request, file: UploadFile = File(...)):
+async def analyze_image(
+    request: Request,
+    background_tasks: BackgroundTasks,
+    file: UploadFile = File(...)
+):
     """Analyze chest X-ray image"""
     if not file.content_type.startswith('image/'):
-        raise HTTPException(400, "Only image files are accepted")
+        raise HTTPException(400, "Only image files accepted")
     
     if file.size > MAX_FILE_SIZE:
-        raise HTTPException(413, f"Max file size is {MAX_FILE_SIZE//(1024*1024)}MB")
+        raise HTTPException(413, f"File too large (max {MAX_FILE_SIZE//1024//1024}MB)")
     
     try:
         contents = await file.read()
@@ -210,56 +216,49 @@ async def analyze_image(request: Request, file: UploadFile = File(...)):
         
         # Store heatmap with session ID
         session_id = str(uuid.uuid4())
-        img_bytes = io.BytesIO()
-        heatmap.save(img_bytes, format='PNG')
+        heatmap_path = heatmap_dir / f"{session_id}.png"
+        heatmap.save(heatmap_path)
         
-        heatmap_store[session_id] = {
-            'image': img_bytes.getvalue(),
-            'timestamp': datetime.now()
-        }
-        cleanup_expired_heatmaps()
+        # Add cleanup task
+        background_tasks.add_task(cleanup_old_heatmaps)
+        
+        # Generate HTTPS URL
+        heatmap_url = f"https://{request.url.hostname}/static/heatmap/{session_id}.png"
         
         return {
             "session_id": session_id,
             "predictions": decoded[0],
-            "heatmap_url": f"{request.base_url}static/heatmap/{session_id}"
+            "heatmap_url": heatmap_url
         }
     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"""
-    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",
-        headers={"Cache-Control": "max-age=300"}
-    )
+        raise HTTPException(500, f"Analysis failed: {str(e)}")
 
 @app.get("/model/info")
 async def model_info():
     """Get model metadata"""
     return {
         "model_type": "DenseNet121",
-        "input_size": "540x540",
-        "classes": len(class_names),
+        "input_shape": str(model.input_shape),
         "output_shape": str(model.output_shape),
+        "classes": len(class_names),
         "gradcam_layer": layer_name,
         "rate_limit": "5 requests/minute"
     }
 
 @app.exception_handler(HTTPException)
-async def http_handler(request: Request, exc: HTTPException):
+async def http_exception_handler(request, exc):
     return JSONResponse(
         status_code=exc.status_code,
         content={"error": exc.detail}
     )
 
 @app.exception_handler(Exception)
-async def generic_handler(request: Request, exc: Exception):
+async def generic_exception_handler(request, exc):
     return JSONResponse(
         status_code=500,
         content={"error": "Internal server error"}
-    )
+    )
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=PORT)