Create main.py

main.py

@@ -0,0 +1,168 @@
+import torch
+import base64
+import io
+import time
+import numpy as np
+from fastapi import FastAPI, File, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
+from torchvision import models, transforms
+from PIL import Image
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+
+app = FastAPI()
+
+# Enable CORS so your React app can talk to this backend
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"], # Allow all origins for the demo
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# --- 1. CONFIGURATION ---
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+CLASSES = ['Gray Leaf Spot', 'Healthy']
+
+# Define paths to your uploaded weights
+# Upload these files to your HF Space manually via "Files" tab
+MODEL_PATHS = {
+    "resnet_base": "models/ResNet50_Aug_False.pth",
+    "resnet_aug": "models/ResNet50_Aug_True.pth",
+    "effnet_base": "models/EfficientNet_Aug_False.pth",
+    "effnet_aug": "models/EfficientNet_Aug_True.pth"
+}
+
+# --- 2. LOAD MODELS ---
+loaded_models = {}
+
+def load_architecture(model_name, num_classes=2):
+    """Rebuilds the architecture to match your training"""
+    if "resnet" in model_name:
+        model = models.resnet50(weights=None)
+        model.fc = torch.nn.Linear(model.fc.in_features, num_classes)
+        # Target layer for Grad-CAM in ResNet
+        target_layer = model.layer4[-1]
+    else:
+        model = models.efficientnet_b0(weights=None)
+        model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, num_classes)
+        # Target layer for Grad-CAM in EfficientNet
+        target_layer = model.features[-1]
+    return model, target_layer
+
+print("Loading models... this might take a minute...")
+for key, path in MODEL_PATHS.items():
+    try:
+        # Create architecture
+        model, layer = load_architecture(key)
+        # Load weights (Ensure you upload the files!)
+        # If testing without weights, comment out the next line
+        state_dict = torch.load(path, map_location=DEVICE)
+        model.load_state_dict(state_dict)
+        
+        model.to(DEVICE)
+        model.eval()
+        
+        # Initialize Grad-CAM for this model
+        cam = GradCAM(model=model, target_layers=[layer])
+        
+        loaded_models[key] = {"model": model, "cam": cam}
+        print(f"Loaded {key}")
+    except Exception as e:
+        print(f"Error loading {key}: {e}")
+        # Placeholder for demo if weights are missing
+        loaded_models[key] = None
+
+# --- 3. UTILITIES ---
+transform = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+def run_inference_and_gradcam(key, image_tensor, original_image_np):
+    """Runs prediction and generates heatmap for a single model"""
+    item = loaded_models[key]
+    if item is None:
+        return None
+
+    model = item["model"]
+    cam = item["cam"]
+    
+    start_time = time.time()
+    
+    # 1. Prediction
+    with torch.no_grad():
+        outputs = model(image_tensor)
+        probs = torch.nn.functional.softmax(outputs, dim=1)
+        conf, pred_idx = torch.max(probs, 1)
+        
+    inference_time = (time.time() - start_time) * 1000 # ms
+    
+    # 2. Grad-CAM
+    # We need gradients, so we run cam() which handles the forward/backward internally
+    grayscale_cam = cam(input_tensor=image_tensor, targets=None)[0, :]
+    visualization = show_cam_on_image(original_image_np, grayscale_cam, use_rgb=True)
+    
+    # Convert Grad-CAM numpy to Base64 String for frontend
+    pil_img = Image.fromarray(visualization)
+    buff = io.BytesIO()
+    pil_img.save(buff, format="JPEG")
+    img_str = base64.b64encode(buff.getvalue()).decode("utf-8")
+    
+    return {
+        "label": CLASSES[pred_idx.item()],
+        "confidence": float(conf.item()),
+        "time": f"{inference_time:.2f}ms",
+        "heatmap": f"data:image/jpeg;base64,{img_str}"
+    }
+
+# --- 4. API ENDPOINT ---
+@app.post("/analyze")
+async def analyze_leaf(file: UploadFile = File(...)):
+    # Read Image
+    contents = await file.read()
+    image = Image.open(io.BytesIO(contents)).convert("RGB")
+    
+    # Prepare Inputs
+    tensor = transform(image).unsqueeze(0).to(DEVICE)
+    
+    # For Grad-CAM visualization, we need a normalized float numpy array (0-1)
+    # Resize original image to 256x256 to match tensor
+    img_resized = image.resize((256, 256))
+    img_np = np.array(img_resized, dtype=np.float32) / 255.0
+
+    results = []
+    
+    # Process all 4 models
+    # Mapping frontend IDs to backend keys
+    definitions = [
+        {"id": 1, "key": "resnet_base", "name": "ResNet50 Base"},
+        {"id": 2, "key": "resnet_aug", "name": "ResNet50 Aug"},
+        {"id": 3, "key": "effnet_base", "name": "EffNet Base"},
+        {"id": 4, "key": "effnet_aug", "name": "EffNet Aug"},
+    ]
+    
+    for definition in definitions:
+        data = run_inference_and_gradcam(definition["key"], tensor, img_np)
+        if data:
+            results.append({
+                "id": definition["id"],
+                **data
+            })
+        else:
+            # Fallback if model failed to load
+            results.append({
+                "id": definition["id"],
+                "label": "Error",
+                "confidence": 0.0,
+                "time": "0ms",
+                "heatmap": ""
+            })
+            
+    return results
+
+@app.get("/")
+def home():
+    return {"message": "Maize Ablation Backend is Running"}