Update app.py

app.py

@@ -1,7 +1,6 @@
 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
+from fastapi.responses import JSONResponse
 from slowapi import Limiter
 from slowapi.util import get_remote_address
 import tensorflow as tf
@@ -17,33 +16,22 @@ import cv2
 import io
 import uuid
 from datetime import datetime, timedelta
-from pathlib import Path
-import os
+import base64
 
 # Configuration
 MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB
-HEATMAP_EXPIRY = 300  # 5 minutes in seconds
 PORT = 7860
-HEATMAP_DIR = "/tmp/heatmaps"  # Changed to writable /tmp directory
 
-# Initialize FastAPI
 app = FastAPI(
     title="ChexNet Medical Imaging API",
     description="API for chest X-ray analysis with Grad-CAM visualization",
-    version="4.1.0"
+    version="5.0.0"
 )
 
 # Rate limiter setup
 limiter = Limiter(key_func=get_remote_address)
 app.state.limiter = limiter
 
-# Create heatmap directory (in /tmp which is writable)
-heatmap_dir = Path(HEATMAP_DIR)
-heatmap_dir.mkdir(parents=True, exist_ok=True)
-
-# Mount static files from /tmp
-app.mount("/static/heatmap", StaticFiles(directory=HEATMAP_DIR), name="heatmaps")
-
 # CORS configuration
 app.add_middleware(
     CORSMiddleware,
@@ -94,16 +82,6 @@ except Exception as e:
     print(f"❌ Model loading failed: {e}")
     raise
 
-async def cleanup_old_heatmaps():
-    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()
-            except Exception as e:
-                print(f"Error deleting {file.name}: {e}")
-
 def generate_gradcam(img):
     img_array = img_to_array(img)
     img_array = np.expand_dims(img_array, axis=0)
@@ -149,7 +127,6 @@ def preprocess_image(file_bytes):
 @limiter.limit("5/minute")
 async def analyze_image(
     request: Request,
-    background_tasks: BackgroundTasks,
     file: UploadFile = File(...)
 ):
     if not file.content_type.startswith('image/'):
@@ -171,16 +148,15 @@ async def analyze_image(
         
         heatmap = generate_gradcam(img)
         
-        session_id = str(uuid.uuid4())
-        heatmap_path = heatmap_dir / f"{session_id}.png"
-        heatmap.save(heatmap_path)
-        
-        background_tasks.add_task(cleanup_old_heatmaps)
+        # Convert heatmap to base64 instead of saving to file
+        img_byte_arr = io.BytesIO()
+        heatmap.save(img_byte_arr, format='PNG')
+        heatmap_base64 = base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
         
         return {
-            "session_id": session_id,
             "predictions": decoded[0],
-            "heatmap_url": f"https://{request.url.hostname}/static/heatmap/{session_id}.png"
+            "heatmap_image": heatmap_base64,
+            "heatmap_format": "base64 encoded PNG"
         }
     except Exception as e:
         raise HTTPException(500, f"Analysis failed: {str(e)}")
@@ -189,8 +165,7 @@ async def analyze_image(
 async def health_check():
     return {
         "status": "healthy",
-        "timestamp": datetime.now().isoformat(),
-        "heatmap_files": len(list(heatmap_dir.glob("*.png")))
+        "timestamp": datetime.now().isoformat()
     }
 
 if __name__ == "__main__":