Update app.py

app.py

@@ -9,30 +9,398 @@ import pickle
 import numpy as np
 import cv2
 from PIL import Image
-from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi import FastAPI, File, UploadFile, HTTPException, status
+from pydantic import BaseModel
+from typing import List, Dict, Any
 
-# --- START OF FIX ---
-# We must import torch FIRST.
+# --- IMPORT ORDER FIX ---
+# 1. Import torch FIRST
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-# NOW, we apply the global patch.
-# This fixes the "Custom classes... not available" error
-# by putting the patched Path object in the global namespace.
-import pathlib
-pathlib.PosixPath = pathlib.Path
+# 2. REMOVE the pathlib patch
+# The global patch was breaking matplotlib, which fastai imports.
+# import pathlib
+# pathlib.PosixPath = pathlib.Path
 
-# FINALLY, we import fastai.
+# 3. Import fastai LAST
 from fastai.vision.all import load_learner, PILImage
-# --- END OF FIX ---
+# --- END IMPORT ORDER FIX ---
 
 
 # =======================
 # CONFIG
 # =======================
-# ... (rest of the file is unchanged) ...
-# ...
+class Config:
+    IMG_SIZE_CLF = 224
+    CAM_PERCENTILE = 75
+    MIN_AREA_RATIO = 0.01
+
+cfg = Config()
+
+# =======================
+# GRAD-CAM IMPLEMENTATION
+# =======================
+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
+                    last_conv_name = name
+        
+        if last_conv is None:
+            raise RuntimeError("No Conv2d layer found in model")
+        
+        return last_conv
+    
+    def compute(self, img_path, target_class_idx):
+        """Compute Grad-CAM for a single 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_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)
+            
+            xb = img_tensor.to(self.device)
+            xb = (xb - mean) / std
+            xb = xb.requires_grad_(True)
+            
+            # Hook storage
+            activations_list = []
+            gradients_list = []
+            
+            def save_activation(module, input, output):
+                activations_list.clear()
+                activations_list.append(output.detach().clone())
+            
+            def save_gradient(module, grad_in, grad_out):
+                gradients_list.clear()
+                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
+            cam_resized = F.interpolate(
+                cam_map.unsqueeze(0).unsqueeze(0),
+                size=(orig_h, orig_w),
+                mode='bilinear',
+                align_corners=False
+            ).squeeze()
+            
+            # Normalize
+            cam_min = cam_resized.min()
+            cam_max = cam_resized.max()
+            
+            if cam_max - cam_min > 1e-8:
+                cam_normalized = (cam_resized - cam_min) / (cam_max - cam_min)
+            else:
+                cam_normalized = torch.zeros_like(cam_resized)
+            
+            # Cleanup
+            fwd_handle.remove()
+            bwd_handle.remove()
+            self.model.zero_grad()
+            
+            return cam_normalized.clamp(0, 1).detach().cpu()
+        
+        except Exception as e:
+            print(f"⚠ Grad-CAM error: {e}", file=sys.stderr)
+            return None
+
+# =======================
+# BBOX GENERATION
+# =======================
+def cam_to_multiscale_bboxes(cam, img_w, img_h):
+    """Generate multiple bboxes at different thresholds."""
+    
+    if cam is None:
+        return []
+    
+    cam_np = cam.numpy() if isinstance(cam, torch.Tensor) else cam
+    cam_np = (cam_np * 255).astype(np.uint8)
+    
+    boxes = []
+    img_area = img_w * img_h
+    
+    # Try multiple thresholds
+    percentiles = [60, 75, 85]
+    seen_boxes = set()
+    
+    for percentile in percentiles:
+        non_zero_mask = cam_np > 0
+        if not np.any(non_zero_mask):
+            continue
+        
+        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)
+        
+        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        
+        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
+
+def apply_nms(boxes, iou_threshold=0.5):
+    """Non-Maximum Suppression."""
+    if len(boxes) == 0:
+        return []
+    
+    boxes = np.array(boxes)
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 0] + boxes[:, 2]
+    y2 = boxes[:, 1] + boxes[:, 3]
+    scores = boxes[:, 4]
+    
+    areas = (x2 - x1) * (y2 - y1)
+    order = scores.argsort()[::-1]
+    
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+        
+        w = np.maximum(0.0, xx2 - xx1)
+        h = np.maximum(0.0, yy2 - yy1)
+        inter = w * h
+        
+        iou = inter / (areas[i] + areas[order[1:]] - inter)
+        
+        inds = np.where(iou <= iou_threshold)[0]
+        order = order[inds + 1]
+    
+    return boxes[keep].tolist()
+
+def filter_contained_boxes(boxes, tolerance=10):
+    """Filter out boxes that are contained within larger boxes with tolerance."""
+    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 = []
+    
+    for box in boxes_sorted:
+        contained = False
+        for larger_box in filtered:
+            if is_contained(box, larger_box, tolerance):
+                contained = True
+                break
+        if not contained:
+            filtered.append(box)
+    
+    return filtered
+
+def is_contained(small_box, large_box, tolerance):
+    """Check if small_box is contained within large_box with tolerance."""
+    sx, sy, sw, sh = small_box[:4]
+    lx, ly, lw, lh = large_box[:4]
+    
+    return (sx >= lx - tolerance and
+            sy >= ly - tolerance and
+            sx + sw <= lx + lw + tolerance and
+            sy + sh <= ly + lh + tolerance)
+
+# =======================
+# MODIFIED MAIN INFERENCE
+# =======================
+def run_inference(image_path, 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_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
+    
+    with torch.no_grad():
+        output = learn.model(xb.to(device))
+    
+    pred_idx = output.argmax(dim=1).item()
+    predicted_class = class_names[pred_idx]
+    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
+    
+    # Generate Grad-CAM
+    gradcam = GradCAM(learn)
+    cam = gradcam.compute(image_path, 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
+            'boundingBox': {
+                'x': int(x),
+                'y': int(y),
+                'width': int(w),
+                'height': int(h)
+            },
+            'classId': pred_idx
+        })
+    
+    # If no boxes found, return full image as bbox
+    if len(detections) == 0:
+        detections.append({
+            'diseaseName': predicted_class,
+            'confidence': confidence,
+            'boundingBox': {
+                'x': 0,
+                'y': 0,
+                'width': img_w,
+                'height': img_h
+            },
+            'classId': pred_idx
+        })
+    
+    return detections
+
 # =======================
 # FASTAPI SERVER
 # =======================
@@ -40,7 +408,6 @@ from fastai.vision.all import load_learner, PILImage
 # Store model in a global cache
 class ModelCache:
     learn = None
-    class_names = None
 
 model_cache = ModelCache()
 
@@ -53,11 +420,12 @@ async def lifespan(app: FastAPI):
         print(f"FATAL: Model file not found at {model_path}", file=sys.stderr)
     else:
         try:
-            # We NO LONGER need the patch here, it's global now.
+            # We have REMOVED the pathlib patch.
+            # If this fails, the model was saved with a patch and
+            # this is a more complex problem.
             
             # Force CPU loading
             model_cache.learn = load_learner(model_path, cpu=True)
-            model_cache.class_names = model_cache.learn.dls.vocab
             print("Learner loaded successfully.")
 
         except Exception as e:
@@ -65,7 +433,70 @@ async def lifespan(app: FastAPI):
     yield
     # Clear model from memory on shutdown
     model_cache.learn = None
-    model_cache.class_names = None
     print("Model cache cleared.")
 
-# ... (rest of the file is unchanged) ...
+# 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]
+
+# Initialize FastAPI app with the lifespan event handler
+app = FastAPI(lifespan=lifespan)
+
+@app.get("/")
+def read_root():
+    """Root endpoint for health check."""
+    return {"status": "ok", "model_loaded": model_cache.learn is not None}
+
+@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."
+        )
+
+    # Define a temporary path to save the uploaded image
+    # Using /tmp/ is standard for temporary files in Linux containers
+    temp_image_path = f"/tmp/{file.filename}"
+
+    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)
+        
+        # Run inference using the saved file path
+        detections = run_inference(temp_image_path, model_cache.learn)
+        
+        # Return the formatted detections
+        return {"detections": detections}
+    
+    except Exception as e:
+        print(f"Error during prediction: {e}", file=sys.stderr)
+        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)
+
+# Note: The `if __name__ == "__main__":` block is removed.
+# Uvicorn will run this "app" object.