Update app.py

app.py

@@ -1,106 +1,55 @@
-import sys
+"""
+FastAPI Disease Detection Service with Grad-CAM
+Usage: uvicorn app:app --host 0.0.0.0 --port 8000
+"""
+
+import io
 import json
 import warnings
-import os
-import aiofiles
-from contextlib import asynccontextmanager
-from pathlib import Path
-import pickle
-import platform
-import io
-
 warnings.filterwarnings('ignore')
 
+
+
+from typing import List, Optional
+from fastapi import FastAPI, File, UploadFile, HTTPException, Query
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import uvicorn
+
 import numpy as np
 import cv2
+from pathlib import Path
 from PIL import Image
-from fastapi import FastAPI, File, UploadFile, HTTPException, status
-from pydantic import BaseModel
-from typing import List, Dict, Any
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from fastai.vision.all import PILImage, Learner
+from fastai.vision.all import load_learner, PILImage
 
 # =======================
-# CUSTOM UNPICKLER FOR CROSS-PLATFORM COMPATIBILITY
+# PYDANTIC MODELS
 # =======================
-class CrossPlatformUnpickler(pickle.Unpickler):
-    """Custom unpickler to handle Windows paths on Linux."""
-    
-    def find_class(self, module, name):
-        """Override to handle pathlib classes."""
-        if module == 'pathlib':
-            # Map both PosixPath and WindowsPath to the generic Path
-            if name in ('PosixPath', 'WindowsPath'):
-                import pathlib
-                return pathlib.Path
-        return super().find_class(module, name)
+class BoundingBox(BaseModel):
+    x: int
+    y: int
+    width: int
+    height: int
 
-def load_model_cross_platform(model_path):
-    """Load fastai model with cross-platform pathlib compatibility."""
-    print(f"Attempting to load model from: {model_path}")
-    
-    try:
-        # Read the pickle file
-        with open(model_path, 'rb') as f:
-            pkl_data = f.read()
-        
-        # Use custom unpickler
-        print("Unpickling with custom unpickler...")
-        learner = CrossPlatformUnpickler(io.BytesIO(pkl_data)).load()
-        
-        # If it's already a Learner object, return it
-        if isinstance(learner, Learner):
-            print("✓ Learner unpickled directly")
-            learner.dls.cpu()
-            return learner
-        
-        # If it's a dict or other structure, try to extract the learner
-        print(f"Unpickled object type: {type(learner)}")
-        
-        # fastai sometimes wraps the learner in a dict
-        if isinstance(learner, dict):
-            if 'learner' in learner:
-                learner = learner['learner']
-            elif 'model' in learner:
-                print("Found model in dict, attempting to reconstruct learner...")
-                # This is trickier - you may need to reconstruct the Learner
-                raise ValueError("Model dict format not directly supported. Please re-export your model.")
-        
-        if isinstance(learner, Learner):
-            learner.dls.cpu()
-            return learner
-        else:
-            raise ValueError(f"Unexpected unpickled type: {type(learner)}")
-            
-    except Exception as e:
-        print(f"Custom unpickler failed: {e}")
-        print("Attempting fallback with pathlib patch...")
-        
-        # Fallback: Try with pathlib patch
-        import pathlib
-        original_posix = getattr(pathlib, 'PosixPath', None)
-        original_windows = getattr(pathlib, 'WindowsPath', None)
-        
-        try:
-            # Patch pathlib
-            if platform.system() != 'Windows':
-                pathlib.WindowsPath = pathlib.Path
-                pathlib.PosixPath = pathlib.Path
-            
-            # Try standard fastai loader
-            from fastai.vision.all import load_learner
-            learner = load_learner(model_path, cpu=True)
-            return learner
-            
-        finally:
-            # Restore pathlib
-            if original_posix is not None:
-                pathlib.PosixPath = original_posix
-            if original_windows is not None:
-                pathlib.WindowsPath = original_windows
+class Detection(BaseModel):
+    diseaseName: str
+    confidence: float
+    boundingBox: BoundingBox
+    classId: int
+
+class InferenceResponse(BaseModel):
+    success: bool
+    detections: List[Detection]
+    message: Optional[str] = None
+
+class HealthResponse(BaseModel):
+    status: str
+    model_loaded: bool
+    device: str
 
 # =======================
 # CONFIG
@@ -109,7 +58,8 @@ class Config:
     IMG_SIZE_CLF = 224
     CAM_PERCENTILE = 75
     MIN_AREA_RATIO = 0.01
-
+    MODEL_PATH = "classifier.pkl"  # Default model path
+    
 cfg = Config()
 
 # =======================
@@ -117,28 +67,25 @@ cfg = Config()
 # =======================
 class GradCAM:
     """Grad-CAM for single image inference."""
-
+    
     def __init__(self, learn):
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.device = device
         self.model = learn.model.to(device).eval()
         self.target_layer = self._find_target_layer()
-
+    
     def _find_target_layer(self):
         """Find last spatial conv layer (not 1x1 convolutions)."""
         last_conv = None
         last_conv_name = None
-
-        # Iterate through all modules
+        
         for name, m in self.model.named_modules():
             if isinstance(m, nn.Conv2d):
-                # Skip 1x1 convolutions (classifier heads)
                 if m.kernel_size != (1, 1):
                     last_conv = m
                     last_conv_name = name
         
         if last_conv is None:
-            # Fallback: try to find ANY conv layer
             for name, m in self.model.named_modules():
                 if isinstance(m, nn.Conv2d):
                     last_conv = m
@@ -149,28 +96,22 @@ class GradCAM:
         
         return last_conv
     
-    def compute(self, img_path, target_class_idx):
-        """Compute Grad-CAM for a single image."""
+    def compute(self, img_pil, target_class_idx):
+        """Compute Grad-CAM for a PIL image."""
         
         try:
-            # Load and preprocess image
-            img = PILImage.create(img_path)
-            img_np = np.array(img.resize((cfg.IMG_SIZE_CLF, cfg.IMG_SIZE_CLF)))
+            img_np = np.array(img_pil.resize((cfg.IMG_SIZE_CLF, cfg.IMG_SIZE_CLF)))
             img_tensor = torch.from_numpy(img_np).float() / 255.0
             
-            # Handle grayscale
             if img_tensor.ndim == 2:
                 img_tensor = img_tensor.unsqueeze(0).repeat(3, 1, 1)
             elif img_tensor.ndim == 3:
                 img_tensor = img_tensor.permute(2, 0, 1)
-                # Ensure 3 channels
                 if img_tensor.shape[0] == 1:
                     img_tensor = img_tensor.repeat(3, 1, 1)
             
-            # Add batch dimension
             img_tensor = img_tensor.unsqueeze(0)
             
-            # ImageNet normalization
             mean = torch.tensor([0.485, 0.456, 0.406], device=self.device).view(1, 3, 1, 1)
             std = torch.tensor([0.229, 0.224, 0.225], device=self.device).view(1, 3, 1, 1)
             
@@ -178,7 +119,6 @@ class GradCAM:
             xb = (xb - mean) / std
             xb = xb.requires_grad_(True)
             
-            # Hook storage
             activations_list = []
             gradients_list = []
             
@@ -191,41 +131,30 @@ class GradCAM:
                 if grad_out[0] is not None:
                     gradients_list.append(grad_out[0].detach().clone())
             
-            # Register hooks
             fwd_handle = self.target_layer.register_forward_hook(save_activation)
             bwd_handle = self.target_layer.register_full_backward_hook(save_gradient)
             
-            # Forward pass
             self.model.zero_grad()
             with torch.set_grad_enabled(True):
                 output = self.model(xb)
             
-            # Check activations
             if len(activations_list) == 0:
-                print(f"⚠ Warning: Forward hook didn't fire", file=sys.stderr)
                 return None
             
-            # Backward pass
             target_score = output[0, target_class_idx]
             target_score.backward()
             
-            # Check gradients
             if len(gradients_list) == 0:
-                print(f"⚠ Warning: Backward hook didn't fire", file=sys.stderr)
                 return None
             
-            # Get activations and gradients
             acts = activations_list[0].to(self.device)
             grads = gradients_list[0].to(self.device)
             
-            # Compute CAM
             weights = grads.mean(dim=[2, 3], keepdim=True)
             cam_map = (weights * acts).sum(dim=1).squeeze(0)
             cam_map = F.relu(cam_map)
             
-            # Resize to original size
-            orig_img = Image.open(img_path)
-            orig_w, orig_h = orig_img.size
+            orig_w, orig_h = img_pil.size
             cam_resized = F.interpolate(
                 cam_map.unsqueeze(0).unsqueeze(0),
                 size=(orig_h, orig_w),
@@ -233,7 +162,6 @@ class GradCAM:
                 align_corners=False
             ).squeeze()
             
-            # Normalize
             cam_min = cam_resized.min()
             cam_max = cam_resized.max()
             
@@ -242,7 +170,6 @@ class GradCAM:
             else:
                 cam_normalized = torch.zeros_like(cam_resized)
             
-            # Cleanup
             fwd_handle.remove()
             bwd_handle.remove()
             self.model.zero_grad()
@@ -250,7 +177,7 @@ class GradCAM:
             return cam_normalized.clamp(0, 1).detach().cpu()
         
         except Exception as e:
-            print(f"⚠ Grad-CAM error: {e}", file=sys.stderr)
+            print(f"Grad-CAM error: {e}")
             return None
 
 # =======================
@@ -268,7 +195,6 @@ def cam_to_multiscale_bboxes(cam, img_w, img_h):
     boxes = []
     img_area = img_w * img_h
     
-    # Try multiple thresholds
     percentiles = [60, 75, 85]
     seen_boxes = set()
     
@@ -280,7 +206,6 @@ def cam_to_multiscale_bboxes(cam, img_w, img_h):
         thresh_val = np.percentile(cam_np[non_zero_mask], percentile)
         _, thresh = cv2.threshold(cam_np, int(thresh_val), 255, cv2.THRESH_BINARY)
         
-        # Morphological cleanup
         kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
         thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=1)
         thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=1)
@@ -290,31 +215,25 @@ def cam_to_multiscale_bboxes(cam, img_w, img_h):
         for cnt in contours:
             area = cv2.contourArea(cnt)
             
-            # Dynamic min_area based on threshold
             min_area_ratio = 0.005 if percentile == 60 else 0.01
             min_area = min_area_ratio * img_area
             
             if area > min_area:
                 x, y, w, h = cv2.boundingRect(cnt)
                 
-                # Filter tiny boxes
                 if w < 10 or h < 10:
                     continue
                 
-                # Avoid duplicates
                 box_key = (x // 5, y // 5, w // 5, h // 5)
                 if box_key not in seen_boxes:
                     seen_boxes.add(box_key)
                     
-                    # Confidence based on area and threshold
                     conf = (area / img_area) * (percentile / 100.0)
                     boxes.append([x, y, w, h, min(conf, 1.0)])
     
-    # Apply NMS
     if len(boxes) > 1:
         boxes = apply_nms(boxes, iou_threshold=0.5)
     
-    # Filter contained boxes
     boxes = filter_contained_boxes(boxes, tolerance=10)
     
     return boxes
@@ -356,11 +275,10 @@ def apply_nms(boxes, iou_threshold=0.5):
     return boxes[keep].tolist()
 
 def filter_contained_boxes(boxes, tolerance=10):
-    """Filter out boxes that are contained within larger boxes with tolerance."""
+    """Filter out boxes that are contained within larger boxes."""
     if len(boxes) <= 1:
         return boxes
     
-    # Sort by area descending (larger first)
     boxes_sorted = sorted(boxes, key=lambda b: b[2] * b[3], reverse=True)
     filtered = []
     
@@ -376,7 +294,7 @@ def filter_contained_boxes(boxes, tolerance=10):
     return filtered
 
 def is_contained(small_box, large_box, tolerance):
-    """Check if small_box is contained within large_box with tolerance."""
+    """Check if small_box is contained within large_box."""
     sx, sy, sw, sh = small_box[:4]
     lx, ly, lw, lh = large_box[:4]
     
@@ -386,36 +304,28 @@ def is_contained(small_box, large_box, tolerance):
             sy + sh <= ly + lh + tolerance)
 
 # =======================
-# INFERENCE
+# INFERENCE LOGIC
 # =======================
-def run_inference(image_path, learn):
+def run_inference(img_pil, learn):
     """Run inference using classifier + Grad-CAM."""
     
-    # Get class names from the loaded learner
     class_names = learn.dls.vocab
     
-    # Get prediction
-    img = PILImage.create(image_path)
-    
     # Manual preprocessing
-    img_np = np.array(img.resize((cfg.IMG_SIZE_CLF, cfg.IMG_SIZE_CLF)))
+    img_np = np.array(img_pil.resize((cfg.IMG_SIZE_CLF, cfg.IMG_SIZE_CLF)))
     img_tensor = torch.from_numpy(img_np).float() / 255.0
     
-    # Handle grayscale
     if img_tensor.ndim == 2:
         img_tensor = img_tensor.unsqueeze(0).repeat(3, 1, 1)
     elif img_tensor.ndim == 3:
         img_tensor = img_tensor.permute(2, 0, 1)
-        # Ensure 3 channels
         if img_tensor.shape[0] == 1:
             img_tensor = img_tensor.repeat(3, 1, 1)
     
-    # Add batch dimension
     img_tensor = img_tensor.unsqueeze(0)
     
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     
-    # ImageNet normalization
     mean = torch.tensor([0.485, 0.456, 0.406], device=device).view(1, 3, 1, 1)
     std = torch.tensor([0.229, 0.224, 0.225], device=device).view(1, 3, 1, 1)
     xb = (img_tensor - mean) / std
@@ -428,27 +338,20 @@ def run_inference(image_path, learn):
     probs = F.softmax(output, dim=1).squeeze(0)
     confidence = probs[pred_idx].item()
     
-    # Get image dimensions
-    orig_img = Image.open(image_path)
-    img_w, img_h = orig_img.size
+    img_w, img_h = img_pil.size
     
-    # Generate Grad-CAM
     gradcam = GradCAM(learn)
-    cam = gradcam.compute(image_path, pred_idx)
+    cam = gradcam.compute(img_pil, pred_idx)
     
-    # Generate bounding boxes
     boxes = cam_to_multiscale_bboxes(cam, img_w, img_h)
-    
-    # Filter overlapping boxes
     boxes = filter_contained_boxes(boxes, tolerance=10)
     
-    # Format detections
     detections = []
     for box in boxes:
         x, y, w, h, conf = box
         detections.append({
             'diseaseName': predicted_class,
-            'confidence': float(conf * confidence),  # Combined confidence
+            'confidence': float(conf * confidence),
             'boundingBox': {
                 'x': int(x),
                 'y': int(y),
@@ -458,7 +361,6 @@ def run_inference(image_path, learn):
             'classId': pred_idx
         })
     
-    # If no boxes found, return full image as bbox
     if len(detections) == 0:
         detections.append({
             'diseaseName': predicted_class,
@@ -475,98 +377,148 @@ def run_inference(image_path, learn):
     return detections
 
 # =======================
-# FASTAPI SERVER
+# FASTAPI APP
 # =======================
+app = FastAPI(
+    title="Disease Detection API",
+    description="AI-powered disease detection service with Grad-CAM visualization",
+    version="1.0.0"
+)
 
-# Store model in a global cache
-class ModelCache:
-    learn = None
-
-model_cache = ModelCache()
-
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    # Load the model on startup
-    print("Loading Fastai learner...")
-    model_path = "classifier.pkl"
-    if not Path(model_path).exists():
-        print(f"FATAL: Model file not found at {model_path}", file=sys.stderr)
-    else:
-        try:
-            # Use our safe cross-platform loader
-            model_cache.learn = load_model_cross_platform(model_path)
-            print("✓ Learner loaded successfully.")
-            print(f"✓ Classes: {model_cache.learn.dls.vocab}")
-
-        except Exception as e:
-            print(f"FATAL: Failed to load learner: {e}", file=sys.stderr)
-            import traceback
-            traceback.print_exc()
-    yield
-    # Clear model from memory on shutdown
-    model_cache.learn = None
-    print("Model cache cleared.")
-
-# Define Pydantic models for response
-class BoundingBox(BaseModel):
-    x: int
-    y: int
-    width: int
-    height: int
-
-class Detection(BaseModel):
-    diseaseName: str
-    confidence: float
-    boundingBox: BoundingBox
-    classId: int
-
-class PredictionResponse(BaseModel):
-    detections: List[Detection]
+# CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
-# Initialize FastAPI app with the lifespan event handler
-app = FastAPI(lifespan=lifespan)
+# Global model variable
+model = None
 
-@app.get("/")
-def read_root():
-    """Root endpoint for health check."""
-    return {"status": "ok", "model_loaded": model_cache.learn is not None}
+@app.on_event("startup")
+async def load_model():
+    """Load the classifier model on startup."""
+    global model
+    try:
+        if Path(cfg.MODEL_PATH).exists():
+            model = load_learner(cfg.MODEL_PATH)
+            print(f"✓ Model loaded from {cfg.MODEL_PATH}")
+        else:
+            print(f"⚠ Warning: Model file not found at {cfg.MODEL_PATH}")
+    except Exception as e:
+        print(f"✗ Error loading model: {e}")
 
-@app.post("/predict", response_model=PredictionResponse)
-async def predict(file: UploadFile = File(...)):
-    """Accepts an image, saves it, runs inference, and returns detections."""
-    
-    # Check if model is loaded
-    if model_cache.learn is None:
-        raise HTTPException(
-            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
-            detail="Model is not loaded. Check startup logs."
-        )
+@app.get("/", response_model=HealthResponse)
+async def root():
+    """Root endpoint - health check."""
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    return {
+        "status": "running",
+        "model_loaded": model is not None,
+        "device": device
+    }
 
-    # Define a temporary path to save the uploaded image
-    temp_image_path = f"/tmp/{file.filename}"
+@app.get("/health", response_model=HealthResponse)
+async def health_check():
+    """Health check endpoint."""
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    return {
+        "status": "healthy" if model is not None else "model_not_loaded",
+        "model_loaded": model is not None,
+        "device": device
+    }
 
+@app.post("/predict", response_model=InferenceResponse)
+async def predict(
+    file: UploadFile = File(...),
+    model_path: Optional[str] = Query(None, description="Optional custom model path")
+):
+    """
+    Predict disease from uploaded image.
+    
+    Parameters:
+    - file: Image file (JPG, PNG, etc.)
+    - model_path: Optional custom model path (query parameter)
+    
+    Returns:
+    - JSON with detections including disease name, confidence, and bounding boxes
+    """
+    
+    # Check model
+    current_model = model
+    if model_path:
+        try:
+            if not Path(model_path).exists():
+                raise HTTPException(status_code=400, detail=f"Model not found: {model_path}")
+            current_model = load_learner(model_path)
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"Error loading custom model: {str(e)}")
+    
+    if current_model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded. Please check server logs.")
+    
+    # Validate file type
+    if not file.content_type.startswith('image/'):
+        raise HTTPException(status_code=400, detail="File must be an image")
+    
     try:
-        # Asynchronously save the uploaded file
-        async with aiofiles.open(temp_image_path, 'wb') as out_file:
-            content = await file.read()
-            await out_file.write(content)
+        # Read image
+        contents = await file.read()
+        img_pil = Image.open(io.BytesIO(contents))
+        
+        # Convert RGBA to RGB if needed
+        if img_pil.mode == 'RGBA':
+            img_pil = img_pil.convert('RGB')
         
-        # Run inference using the saved file path
-        detections = run_inference(temp_image_path, model_cache.learn)
+        # Run inference
+        detections = run_inference(img_pil, current_model)
         
-        # Return the formatted detections
-        return {"detections": detections}
+        return {
+            "success": True,
+            "detections": detections,
+            "message": f"Detected {len(detections)} region(s)"
+        }
     
     except Exception as e:
-        print(f"Error during prediction: {e}", file=sys.stderr)
-        import traceback
-        traceback.print_exc()
-        raise HTTPException(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=f"Inference error: {str(e)}"
-        )
-    
-    finally:
-        # Clean up the temporary file
-        if os.path.exists(temp_image_path):
-            os.remove(temp_image_path)
+        raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")
+
+@app.get("/classes")
+async def get_classes():
+    """Get list of disease classes the model can detect."""
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
+    try:
+        classes = list(model.dls.vocab)
+        return {
+            "success": True,
+            "classes": classes,
+            "num_classes": len(classes)
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error retrieving classes: {str(e)}")
+
+# =======================
+# RUN SERVER
+# =======================
+if __name__ == "__main__":
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="Disease Detection API Server")
+    parser.add_argument("--host", default="0.0.0.0", help="Host to bind to")
+    parser.add_argument("--port", type=int, default=8000, help="Port to bind to")
+    parser.add_argument("--model", default="classifier.pkl", help="Path to classifier model")
+    parser.add_argument("--reload", action="store_true", help="Enable auto-reload")
+    
+    args = parser.parse_args()
+    
+    cfg.MODEL_PATH = args.model
+    
+    uvicorn.run(
+        "app:app" if args.reload else app,
+        host=args.host,
+        port=args.port,
+        reload=args.reload
+    )