Update app.py

app.py

@@ -1,190 +1,191 @@
-from fastapi import FastAPI, File, UploadFile, Form, HTTPException
-from fastapi.responses import JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
-import pandas as pd
-import numpy as np
-import torch
-import json
-import io
-import joblib
-import os
-from model import DroughtNetLSTM
-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():
-    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)
-    lstm_dim = 256
-    num_layers = 2
-    dropout = 0.15
-    static_dim = 29  # Number of features in static data
-    staticfc_dim = 16
-    hidden_dim = 256
-    output_size = 6  # Output classes
-    
-    # Initialize model
-    model = DroughtNetLSTM(
-        time_dim=time_dim,
-        lstm_dim=lstm_dim,
-        num_layers=num_layers,
-        dropout=dropout,
-        static_dim=static_dim,
-        staticfc_dim=staticfc_dim,
-        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
-    ))
-    
-    model.to(device)
-    model.eval()
-
-@app.get("/")
-async def root():
-    return {"message": "Welcome to Drought Prediction API. Use /predict endpoint to make predictions."}
-
-@app.post("/predict")
-async def predict(
-    csv_file: UploadFile = File(...),
-    x_static: str = Form(...),
-):
-    try:
-        # Parse x_static from JSON string to list
-        x_static_list = json.loads(x_static)
-        x_static_array = np.array([x_static_list], dtype=np.float32)
-        
-        # Read CSV file
-        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
-        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
-        x_static_norm, x_time_norm = normalize(
-            x_static_array, 
-            x_time_array, 
-            scaler_dict=scaler_dict, 
-            scaler_dict_static=scaler_dict_static
-        )
-        
-        # Convert 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)",
-            2: "Moderate Drought (D2)",
-            3: "Severe Drought (D3)",
-            4: "Extreme Drought (D4)",
-            5: "Exceptional Drought (D5)"
-        }
-        
-        result = {
-            "raw_predictions": predictions,
-            "max_class": {
-                "class": int(np.argmax(predictions)),
-                "label": drought_classes[int(np.argmax(predictions))],
-                "confidence": float(np.max(predictions))
-            },
-            "class_probabilities": {
-                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
+from fastapi import FastAPI, File, UploadFile, Form, HTTPException
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+import pandas as pd
+import numpy as np
+import torch
+import json
+import io
+import joblib
+import os
+from model import DroughtNetLSTM
+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():
+    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)
+    lstm_dim = 256
+    num_layers = 2
+    dropout = 0.15
+    static_dim = 29  # Number of features in static data
+    staticfc_dim = 16
+    hidden_dim = 256
+    output_size = 6  # Output classes
+    print("Khởi tạo dữ scaler hoàn tất")
+    # Initialize model
+    model = DroughtNetLSTM(
+        time_dim=time_dim,
+        lstm_dim=lstm_dim,
+        num_layers=num_layers,
+        dropout=dropout,
+        static_dim=static_dim,
+        staticfc_dim=staticfc_dim,
+        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()
+
+@app.get("/")
+async def root():
+    return {"message": "Welcome to Drought Prediction API. Use /predict endpoint to make predictions."}
+
+@app.post("/predict")
+async def predict(
+    csv_file: UploadFile = File(...),
+    x_static: str = Form(...),
+):
+    try:
+        # Parse x_static from JSON string to list
+        x_static_list = json.loads(x_static)
+        x_static_array = np.array([x_static_list], dtype=np.float32)
+        
+        # Read CSV file
+        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
+        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
+        x_static_norm, x_time_norm = normalize(
+            x_static_array, 
+            x_time_array, 
+            scaler_dict=scaler_dict, 
+            scaler_dict_static=scaler_dict_static
+        )
+        
+        # Convert 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)",
+            2: "Moderate Drought (D2)",
+            3: "Severe Drought (D3)",
+            4: "Extreme Drought (D4)",
+            5: "Exceptional Drought (D5)"
+        }
+        
+        result = {
+            "raw_predictions": predictions,
+            "max_class": {
+                "class": int(np.argmax(predictions)),
+                "label": drought_classes[int(np.argmax(predictions))],
+                "confidence": float(np.max(predictions))
+            },
+            "class_probabilities": {
+                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)