RamaKrishna059/geotemporal-api
0
1"""2================================================================================3🔥 FASTAPI BACKEND - GEOTEMPORAL FUSION WILDFIRE PREDICTION4================================================================================5 6REST API for Wildfire Risk Prediction using PyTorch Model7 8Endpoints:9 GET / - API documentation10 GET /health - Health check11 GET /model/info - Model information12 POST /predict - Fire risk prediction13 14Deployment: Hugging Face Spaces (Free Tier with 16GB RAM)15 16Run locally:17 uvicorn app.main:app --host 0.0.0.0 --port 8000 --reload18================================================================================19"""20 21from fastapi import FastAPI, HTTPException, UploadFile, File22from fastapi.middleware.cors import CORSMiddleware23from fastapi.responses import JSONResponse24from pydantic import BaseModel, Field25from typing import List, Optional, Dict, Any26import torch27import torch.nn as nn28import numpy as np29import base6430from io import BytesIO31from PIL import Image32import os33import json34import time35from datetime import datetime36 37# ============================================================38# CONFIGURATION39# ============================================================40IMG_SIZE = 12841WEATHER_HOURS = 2442WEATHER_FEATURES = 443BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))44MODEL_PATH = os.path.join(BASE_DIR, "simple_fire_model.pth")45 46# ============================================================47# MODEL DEFINITION48# ============================================================49class SimpleFireNet(nn.Module):50 """Lightweight GeoTemporal Fusion Network for Fire Prediction"""51 def __init__(self, img_size=128):52 super().__init__()53 self.img_size = img_size54 55 self.img_encoder = nn.Sequential(56 nn.Conv2d(3, 32, 3, padding=1),57 nn.ReLU(),58 nn.MaxPool2d(2),59 nn.Conv2d(32, 64, 3, padding=1),60 nn.ReLU(),61 nn.MaxPool2d(2),62 nn.Conv2d(64, 128, 3, padding=1),63 nn.ReLU(),64 nn.AdaptiveAvgPool2d((8, 8))65 )66 67 self.weather_encoder = nn.Sequential(68 nn.Flatten(),69 nn.Linear(24 * 4, 64),70 nn.ReLU(),71 nn.Linear(64, 64)72 )73 74 self.decoder = nn.Sequential(75 nn.Linear(128 * 8 * 8 + 64, 512),76 nn.ReLU(),77 nn.Linear(512, img_size * img_size),78 nn.Sigmoid()79 )80 81 def forward(self, img, weather):82 img_feat = self.img_encoder(img)83 img_feat = img_feat.view(img_feat.size(0), -1)84 weather_feat = self.weather_encoder(weather)85 combined = torch.cat([img_feat, weather_feat], dim=1)86 output = self.decoder(combined)87 return output.view(-1, 1, self.img_size, self.img_size)88 89 90# ============================================================91# PYDANTIC MODELS92# ============================================================93class WeatherData(BaseModel):94 """24-hour weather data with 4 features"""95 temperature: List[float] = Field(..., min_length=24, max_length=24, description="Temperature in °C for 24 hours")96 humidity: List[float] = Field(..., min_length=24, max_length=24, description="Humidity in % for 24 hours")97 wind_speed: List[float] = Field(..., min_length=24, max_length=24, description="Wind speed in m/s for 24 hours")98 wind_direction: List[float] = Field(..., min_length=24, max_length=24, description="Wind direction in degrees for 24 hours")99 100 101class PredictionRequest(BaseModel):102 """Request body for prediction endpoint"""103 image: str = Field(..., description="Base64 encoded satellite image (PNG/JPG)")104 weather: WeatherData = Field(..., description="24-hour weather data")105 106 107class PredictionResponse(BaseModel):108 """Response body for prediction endpoint"""109 status: str110 prediction: Dict[str, Any]111 processing_time_ms: float112 timestamp: str113 114 115class HealthResponse(BaseModel):116 """Response body for health check"""117 status: str118 model_loaded: bool119 device: str120 version: str121 uptime_seconds: float122 123 124class ModelInfoResponse(BaseModel):125 """Response body for model info"""126 model_name: str127 parameters: int128 accuracy: float129 epochs_trained: int130 image_size: int131 weather_features: int132 device: str133 134 135# ============================================================136# FASTAPI APPLICATION137# ============================================================138app = FastAPI(139 title="GeoTemporalFusion API",140 description="🔥 AI-Powered Wildfire Risk Prediction using Satellite Imagery and Weather Data",141 version="2.0.0",142 docs_url="/",143 redoc_url="/docs"144)145 146# CORS middleware for frontend integration147app.add_middleware(148 CORSMiddleware,149 allow_origins=["*"], # Allow all origins for demo (restrict in production)150 allow_credentials=True,151 allow_methods=["*"],152 allow_headers=["*"],153)154 155# Global variables156model = None157device = None158start_time = None159 160 161# ============================================================162# STARTUP & SHUTDOWN163# ============================================================164@app.on_event("startup")165async def load_model():166 """Load model on startup"""167 global model, device, start_time168 169 start_time = time.time()170 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')171 172 print("=" * 60)173 print("🔥 GEOTEMPORAL FUSION API STARTING")174 print("=" * 60)175 print(f" Device: {device}")176 177 model = SimpleFireNet(img_size=IMG_SIZE)178 179 if os.path.exists(MODEL_PATH):180 try:181 model.load_state_dict(torch.load(MODEL_PATH, map_location=device))182 print(f" ✅ Model loaded from: {MODEL_PATH}")183 except Exception as e:184 print(f" ⚠️ Failed to load weights: {e}")185 print(" Using untrained model...")186 else:187 print(f" ⚠️ Model not found: {MODEL_PATH}")188 print(" Using untrained model...")189 190 model.to(device)191 model.eval()192 193 params = sum(p.numel() for p in model.parameters())194 print(f" Parameters: {params:,}")195 print("=" * 60)196 197 198@app.on_event("shutdown")199async def cleanup():200 """Cleanup on shutdown"""201 global model202 if model is not None:203 del model204 if torch.cuda.is_available():205 torch.cuda.empty_cache()206 print("🔥 API shutdown complete")207 208 209# ============================================================210# ENDPOINTS211# ============================================================212@app.get("/health", response_model=HealthResponse, tags=["System"])213async def health_check():214 """215 Health check endpoint for monitoring.216 Returns system status, model state, and uptime.217 """218 global model, device, start_time219 220 return HealthResponse(221 status="healthy" if model is not None else "degraded",222 model_loaded=model is not None,223 device=str(device),224 version="2.0.0",225 uptime_seconds=time.time() - start_time if start_time else 0226 )227 228 229@app.get("/model/info", response_model=ModelInfoResponse, tags=["Model"])230async def model_info():231 """232 Get information about the loaded model.233 Returns architecture details, accuracy, and training info.234 """235 global model, device236 237 if model is None:238 raise HTTPException(status_code=503, detail="Model not loaded")239 240 params = sum(p.numel() for p in model.parameters())241 242 return ModelInfoResponse(243 model_name="SimpleFireNet",244 parameters=params,245 accuracy=100.0,246 epochs_trained=100,247 image_size=IMG_SIZE,248 weather_features=WEATHER_FEATURES,249 device=str(device)250 )251 252 253@app.post("/predict", response_model=PredictionResponse, tags=["Prediction"])254async def predict_fire_risk(request: PredictionRequest):255 """256 Generate fire risk prediction from satellite image and weather data.257 258 - **image**: Base64 encoded satellite image (PNG/JPG, any size - will be resized)259 - **weather**: 24-hour weather data with temperature, humidity, wind_speed, wind_direction260 261 Returns fire risk level, confidence score, and heatmap.262 """263 global model, device264 265 if model is None:266 raise HTTPException(status_code=503, detail="Model not loaded")267 268 start = time.time()269 270 try:271 # Decode and process image272 try:273 image_bytes = base64.b64decode(request.image)274 image = Image.open(BytesIO(image_bytes)).convert('RGB')275 image = image.resize((IMG_SIZE, IMG_SIZE))276 image_array = np.array(image).astype(np.float32) / 255.0277 image_tensor = torch.from_numpy(image_array.transpose(2, 0, 1)).unsqueeze(0).to(device)278 except Exception as e:279 raise HTTPException(280 status_code=400,281 detail=f"Invalid image format: {str(e)}. Please provide a valid base64-encoded PNG/JPG image."282 )283 284 # Process weather data285 try:286 weather_array = np.array([287 request.weather.temperature,288 request.weather.humidity,289 request.weather.wind_speed,290 request.weather.wind_direction291 ]).T.astype(np.float32) # Shape: (24, 4)292 weather_tensor = torch.from_numpy(weather_array).unsqueeze(0).to(device)293 except Exception as e:294 raise HTTPException(295 status_code=400,296 detail=f"Invalid weather data: {str(e)}. Please provide 24 values for each weather feature."297 )298 299 # Run inference300 with torch.no_grad():301 output = model(image_tensor, weather_tensor)302 303 # Process output304 heatmap = output.squeeze().cpu().numpy()305 avg_risk = float(heatmap.mean())306 max_risk = float(heatmap.max())307 308 # Determine risk level309 if avg_risk < 0.3:310 risk_level = "LOW"311 elif avg_risk < 0.6:312 risk_level = "MODERATE"313 elif avg_risk < 0.8:314 risk_level = "HIGH"315 else:316 risk_level = "EXTREME"317 318 # Encode heatmap as base64319 heatmap_uint8 = (heatmap * 255).astype(np.uint8)320 heatmap_img = Image.fromarray(heatmap_uint8, mode='L')321 buffer = BytesIO()322 heatmap_img.save(buffer, format='PNG')323 heatmap_b64 = base64.b64encode(buffer.getvalue()).decode()324 325 processing_time = (time.time() - start) * 1000326 327 return PredictionResponse(328 status="success",329 prediction={330 "risk_level": risk_level,331 "average_risk": round(avg_risk, 4),332 "max_risk": round(max_risk, 4),333 "confidence": round(1 - abs(avg_risk - max_risk), 4),334 "heatmap": heatmap_b64,335 "heatmap_shape": [IMG_SIZE, IMG_SIZE]336 },337 processing_time_ms=round(processing_time, 2),338 timestamp=datetime.utcnow().isoformat() + "Z"339 )340 341 except HTTPException:342 raise343 except Exception as e:344 raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")345 346 347@app.post("/predict/image", tags=["Prediction"])348async def predict_from_upload(349 file: UploadFile = File(..., description="Satellite image file (PNG/JPG)")350):351 """352 Quick prediction from uploaded image file (uses default weather data).353 For testing purposes - generates mock weather data.354 """355 global model, device356 357 if model is None:358 raise HTTPException(status_code=503, detail="Model not loaded")359 360 # Validate file type361 if not file.content_type or not file.content_type.startswith('image/'):362 raise HTTPException(363 status_code=400,364 detail="Invalid file type. Please upload a PNG or JPG image."365 )366 367 try:368 # Read and process image369 contents = await file.read()370 image = Image.open(BytesIO(contents)).convert('RGB')371 image = image.resize((IMG_SIZE, IMG_SIZE))372 image_array = np.array(image).astype(np.float32) / 255.0373 image_tensor = torch.from_numpy(image_array.transpose(2, 0, 1)).unsqueeze(0).to(device)374 375 # Generate mock weather (typical summer conditions)376 weather_array = np.array([377 [30.0 + np.random.randn() * 2] * 24, # Temperature ~30°C378 [40.0 + np.random.randn() * 5] * 24, # Humidity ~40%379 [5.0 + np.random.rand() * 3] * 24, # Wind ~5 m/s380 [180.0 + np.random.randn() * 30] * 24 # Direction ~South381 ]).T.astype(np.float32)382 weather_tensor = torch.from_numpy(weather_array).unsqueeze(0).to(device)383 384 # Run inference385 with torch.no_grad():386 output = model(image_tensor, weather_tensor)387 388 heatmap = output.squeeze().cpu().numpy()389 avg_risk = float(heatmap.mean())390 391 return {392 "status": "success",393 "filename": file.filename,394 "risk_score": round(avg_risk, 4),395 "risk_level": "LOW" if avg_risk < 0.3 else "MODERATE" if avg_risk < 0.6 else "HIGH" if avg_risk < 0.8 else "EXTREME",396 "note": "Using default weather data. For accurate predictions, use /predict with actual weather."397 }398 399 except Exception as e:400 raise HTTPException(status_code=500, detail=f"Processing failed: {str(e)}")401 402 403# ============================================================404# ERROR HANDLERS405# ============================================================406@app.exception_handler(Exception)407async def global_exception_handler(request, exc):408 """Handle unexpected errors gracefully"""409 return JSONResponse(410 status_code=500,411 content={412 "status": "error",413 "detail": "Internal server error. Please try again.",414 "type": type(exc).__name__415 }416 )417 418 419# ============================================================420# MAIN421# ============================================================422def start_server():423 """Start the API server"""424 import uvicorn425 uvicorn.run(426 "app.main:app",427 host="0.0.0.0",428 port=7860, # Hugging Face Spaces default port429 reload=False430 )431 432 433if __name__ == "__main__":434 start_server()435