Update app.py

app.py

@@ -16,7 +16,6 @@ import matplotlib.pyplot as plt
 import cv2
 import io
 import uuid
-from typing import Dict
 from datetime import datetime, timedelta
 from pathlib import Path
 import os
@@ -24,25 +23,26 @@ 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
+PORT = 7860
+HEATMAP_DIR = "/tmp/heatmaps"  # Changed to writable /tmp directory
 
-# Initialize FastAPI with rate limiting
+# Initialize FastAPI
 app = FastAPI(
     title="ChexNet Medical Imaging API",
     description="API for chest X-ray analysis with Grad-CAM visualization",
-    version="4.0.0"
+    version="4.1.0"
 )
 
 # Rate limiter setup
 limiter = Limiter(key_func=get_remote_address)
 app.state.limiter = limiter
 
-# Create static/heatmap directory
-heatmap_dir = Path("static/heatmap")
+# Create heatmap directory (in /tmp which is writable)
+heatmap_dir = Path(HEATMAP_DIR)
 heatmap_dir.mkdir(parents=True, exist_ok=True)
 
-# Mount static files
-app.mount("/static", StaticFiles(directory="static"), name="static")
+# Mount static files from /tmp
+app.mount("/static/heatmap", StaticFiles(directory=HEATMAP_DIR), name="heatmaps")
 
 # CORS configuration
 app.add_middleware(
@@ -53,9 +53,6 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Session storage for heatmaps (now using file system)
-heatmap_store: Dict[str, dict] = {}
-
 # Model configuration
 layer_name = 'conv5_block16_concat'
 class_names = [
@@ -65,7 +62,6 @@ class_names = [
 ]
 
 def build_model():
-    """Build DenseNet121 model with correct output shape"""
     base_model = DenseNet121(
         weights=None,
         include_top=False,
@@ -73,20 +69,17 @@ def build_model():
     )
     x = base_model.output
     x = GlobalAveragePooling2D()(x)
-    predictions = Dense(14, activation='sigmoid')(x)  # 14 classes in pretrained weights
+    predictions = Dense(14, activation='sigmoid')(x)
     return Model(inputs=base_model.input, outputs=predictions)
 
 def load_model_with_fallback():
-    """Robust model loading with multiple fallback strategies"""
     try:
-        # Strategy 1: Build model and load weights
         model = build_model()
         model.load_weights('pretrained_model.h5')
         return model
     except Exception as e:
         print(f"Primary loading failed: {e}")
         try:
-            # Strategy 2: Try direct loading
             model = load_model('Densenet.h5', compile=False)
             return model
         except Exception as e:
@@ -97,26 +90,21 @@ def load_model_with_fallback():
 try:
     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
 
 async def cleanup_old_heatmaps():
-    """Delete heatmap files older than HEATMAP_EXPIRY seconds"""
     now = datetime.now()
     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 visualization"""
     img_array = img_to_array(img)
     img_array = np.expand_dims(img_array, axis=0)
     img_array = preprocess_input(img_array)
@@ -146,7 +134,6 @@ def generate_gradcam(img):
     return Image.blend(original_img, heatmap_img, 0.5)
 
 def process_predictions(predictions):
-    """Format model predictions with confidence scores"""
     decoded = []
     for pred in predictions:
         top_indices = np.argsort(pred)[::-1][:len(class_names)]
@@ -154,36 +141,10 @@ def process_predictions(predictions):
     return decoded
 
 def preprocess_image(file_bytes):
-    """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("/")
-async def root():
-    return {
-        "message": "ChexNet API is running",
-        "endpoints": {
-            "docs": "/docs",
-            "health": "/health",
-            "analyze": "POST /analyze"
-        }
-    }
-
-@app.get("/health")
-async def health_check():
-    return {
-        "status": "healthy",
-        "model_loaded": True,
-        "timestamp": datetime.now().isoformat(),
-        "port": PORT,
-        "heatmap_files": len(list(heatmap_dir.glob("*.png")))
-    }
-
-@app.get("/model/classes")
-async def get_class_names():
-    return {"classes": class_names}
-
 @app.post("/analyze")
 @limiter.limit("5/minute")
 async def analyze_image(
@@ -191,7 +152,6 @@ async def analyze_image(
     background_tasks: BackgroundTasks,
     file: UploadFile = File(...)
 ):
-    """Analyze chest X-ray image"""
     if not file.content_type.startswith('image/'):
         raise HTTPException(400, "Only image files accepted")
     
@@ -202,63 +162,37 @@ async def analyze_image(
         contents = await file.read()
         img = preprocess_image(contents)
         
-        # Prepare input tensor
         img_array = img_to_array(img)
         img_array = np.expand_dims(img_array, axis=0)
         img_array = preprocess_input(img_array)
         
-        # Get predictions
         predictions = model.predict(img_array)
         decoded = process_predictions(predictions)
         
-        # Generate Grad-CAM
         heatmap = generate_gradcam(img)
         
-        # Store heatmap with session ID
         session_id = str(uuid.uuid4())
         heatmap_path = heatmap_dir / f"{session_id}.png"
         heatmap.save(heatmap_path)
         
-        # 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": heatmap_url
+            "heatmap_url": f"https://{request.url.hostname}/static/heatmap/{session_id}.png"
         }
     except Exception as e:
         raise HTTPException(500, f"Analysis failed: {str(e)}")
 
-@app.get("/model/info")
-async def model_info():
-    """Get model metadata"""
+@app.get("/health")
+async def health_check():
     return {
-        "model_type": "DenseNet121",
-        "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"
+        "status": "healthy",
+        "timestamp": datetime.now().isoformat(),
+        "heatmap_files": len(list(heatmap_dir.glob("*.png")))
     }
 
-@app.exception_handler(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_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)