Update app.py

app.py

@@ -1,6 +1,8 @@
 from fastapi import FastAPI, File, UploadFile, Form, HTTPException
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
+from contextlib import asynccontextmanager
+
 import pandas as pd
 import numpy as np
 import torch
@@ -13,44 +15,27 @@ from utils import normalize, date_encode, interpolate_nans
 from datetime import datetime
 from typing import List, Optional
 
-app = FastAPI(
-    title="Drought Prediction API",
-    description="API for predicting drought severity based on weather data",
-    version="1.0.0"
-)
-
-# Enable CORS
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Load model and scalers
-@app.on_event("startup")
-async def load_model():
+# Lifespan event handler (thay thế @app.on_event)
+@asynccontextmanager
+async def lifespan(app: FastAPI):
     global model, scaler_dict, scaler_dict_static, device
-    
-    # Set device
+
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    
-    # Load scalers
+
     scaler_dict = joblib.load(os.path.join(os.path.dirname(__file__), "scaler_dict.joblib"))
     scaler_dict_static = joblib.load(os.path.join(os.path.dirname(__file__), "scaler_dict_static.joblib"))
-    
-    # Define model parameters
-    time_dim = 20  # Number of features in time data (18 weather + 2 date encoding)
+    print("Khởi tạo dữ scaler hoàn tất")
+
+    # Define model params
+    time_dim = 20
     lstm_dim = 256
     num_layers = 2
     dropout = 0.15
-    static_dim = 29  # Number of features in static data
+    static_dim = 29
     staticfc_dim = 16
     hidden_dim = 256
-    output_size = 6  # Output classes
-    print("Khởi tạo dữ scaler hoàn tất")
-    # Initialize model
+    output_size = 6
+
     model = DroughtNetLSTM(
         time_dim=time_dim,
         lstm_dim=lstm_dim,
@@ -61,16 +46,31 @@ async def load_model():
         hidden_dim=hidden_dim,
         output_size=output_size
     )
-    
-    # Load model weights
     model.load_state_dict(torch.load(
         os.path.join(os.path.dirname(__file__), "best_macro_f1_model.pt"),
         map_location=device
     ))
-    print("Khởi tạo dữ model hoàn tất")
-    
     model.to(device)
     model.eval()
+    print("Khởi tạo dữ model hoàn tất")
+
+    yield  # Cho phép app chạy
+
+app = FastAPI(
+    title="Drought Prediction API",
+    description="API for predicting drought severity based on weather data",
+    version="1.0.0",
+    lifespan=lifespan
+)
+
+# Enable CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 @app.get("/")
 async def root():
@@ -82,50 +82,44 @@ async def predict(
     x_static: str = Form(...),
 ):
     try:
-        # Parse x_static from JSON string to list
+        # Parse static input
         x_static_list = json.loads(x_static)
         x_static_array = np.array([x_static_list], dtype=np.float32)
-        
-        # Read CSV file
+
+        # Load and process CSV
         content = await csv_file.read()
         df = pd.read_csv(io.StringIO(content.decode('utf-8')))
-        
-        # Prepare time data
         df = prepare_time_data(df)
-        
-        # Get features
+
+        # Feature extraction
         float_cols = [
             'PRECTOTCORR', 'PS', 'QV2M', 'T2M', 'T2MDEW', 'T2MWET', 'T2M_MAX', 'T2M_MIN', 'T2M_RANGE',
             'TS', 'WS10M', 'WS10M_MAX', 'WS10M_MIN', 'WS10M_RANGE',
             'WS50M', 'WS50M_MAX', 'WS50M_MIN', 'WS50M_RANGE',
         ]
-        
         features = float_cols + ['sin_day', 'cos_day']
         x_time_array = df[features].to_numpy(dtype=np.float32)
         x_time_array = np.expand_dims(x_time_array, axis=0)
-        
-        # Normalize data
+
+        # Normalize
         x_static_norm, x_time_norm = normalize(
-            x_static_array, 
-            x_time_array, 
-            scaler_dict=scaler_dict, 
+            x_static_array,
+            x_time_array,
+            scaler_dict=scaler_dict,
             scaler_dict_static=scaler_dict_static
         )
-        
-        # Convert to tensors
+
+        # To tensors
         x_time_tensor = torch.tensor(x_time_norm).float().to(device)
         x_static_tensor = torch.tensor(x_static_norm).float().to(device)
-        
+
         # Predict
         with torch.no_grad():
             output = model(x_time_tensor, x_static_tensor)
-            # Clamp output to [0, 5]
             output = torch.clamp(output, min=0.0, max=5.0)
-            
-        # Convert to list
+
         predictions = output.cpu().numpy().tolist()[0]
-        
-        # Create result with class interpretations
+
         drought_classes = {
             0: "No Drought (D0)",
             1: "Abnormally Dry (D1)",
@@ -134,7 +128,7 @@ async def predict(
             4: "Extreme Drought (D4)",
             5: "Exceptional Drought (D5)"
         }
-        
+
         result = {
             "raw_predictions": predictions,
             "max_class": {
@@ -146,46 +140,38 @@ async def predict(
                 drought_classes[i]: float(predictions[i]) for i in range(len(predictions))
             }
         }
-        
+
         return JSONResponse(content=result)
-    
+
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")
 
+
 def prepare_time_data(df):
-    """
-    Prepare time-series data for the model.
-    """
-    # Ensure we have YEAR and DOY columns
     if 'YEAR' not in df.columns or 'DOY' not in df.columns:
-        # Try to extract from date column if it exists
         if 'date' in df.columns:
             df['date'] = pd.to_datetime(df['date'])
             df['YEAR'] = df['date'].dt.year
             df['DOY'] = df['date'].dt.dayofyear
         else:
             raise ValueError("Input CSV must contain either 'date' column or both 'YEAR' and 'DOY' columns")
-    
-    # Create date column if it doesn't exist
+
     if 'date' not in df.columns:
         df['date'] = pd.to_datetime(df['YEAR'].astype(str) + df['DOY'].astype(str), format="%Y%j")
-    
-    # Apply date encoding to create sin_day and cos_day columns
+
     df[['sin_day', 'cos_day']] = df['date'].apply(lambda d: pd.Series(date_encode(d)))
-    
-    # Handle missing values if any
+
     float_cols = [
         'PRECTOTCORR', 'PS', 'QV2M', 'T2M', 'T2MDEW', 'T2MWET', 'T2M_MAX', 'T2M_MIN', 'T2M_RANGE',
         'TS', 'WS10M', 'WS10M_MAX', 'WS10M_MIN', 'WS10M_RANGE',
         'WS50M', 'WS50M_MAX', 'WS50M_MIN', 'WS50M_RANGE',
     ]
-    
     for col in float_cols:
         if col in df.columns and df[col].isna().any():
             df[col] = interpolate_nans(df[col].values)
-    
+
     return df
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run("app:app", host="0.0.0.0", port=8000)
+    uvicorn.run("app:app", host="0.0.0.0", port=8000)