Upload predict.py

predict.py

@@ -0,0 +1,112 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import pandas as pd
+import joblib
+from flask import Blueprint, request, jsonify
+import os
+
+predict_bp = Blueprint('predict', __name__)
+
+# Lazy load model and preprocessors
+model_bundle = None
+model = None
+preprocessor = None
+nn_model = None
+beach_db = None
+wind_directions = None
+
+class ImprovedTrashPredictorMLP(nn.Module):
+    def __init__(self, input_size, hidden_sizes, output_size, dropout_rate=0.3):
+        super().__init__()
+        layers = []
+        in_features = input_size
+        for hidden_size in hidden_sizes:
+            layers.append(nn.Linear(in_features, hidden_size))
+            layers.append(nn.BatchNorm1d(hidden_size))
+            layers.append(nn.ReLU())
+            layers.append(nn.Dropout(dropout_rate))
+            in_features = hidden_size
+        layers.append(nn.Linear(in_features, output_size))
+        self.model = nn.Sequential(*layers)
+    def forward(self, x):
+        return self.model(x)
+
+def load_bundle(bundle_path):
+    return {
+        "preprocessor": joblib.load(os.path.join(bundle_path, "preprocessor.pkl")),
+        "beach_db": joblib.load(os.path.join(bundle_path, "beach_db.pkl")),
+        "nn_model": joblib.load(os.path.join(bundle_path, "nn_model.pkl")),
+        "wind_directions": joblib.load(os.path.join(bundle_path, "wind_directions.pkl")),
+    }
+
+def lazy_load():
+    global model_bundle, model, preprocessor, nn_model, beach_db, wind_directions
+    if model_bundle is None:
+        bundle_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../saved_models/final_bundle_20250703_112021'))
+        model_bundle = load_bundle(bundle_path)
+        preprocessor = model_bundle["preprocessor"]
+        beach_db = model_bundle["beach_db"]
+        nn_model = model_bundle["nn_model"]
+        wind_directions = model_bundle["wind_directions"]
+        # Model init
+        input_size = preprocessor.transformers_[0][1].steps[0][1].get_feature_names_out().shape[0] + len(preprocessor.transformers_[1][2])
+        model_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../saved_models/final_model_20250703_112021.pth'))
+        model_instance = ImprovedTrashPredictorMLP(
+            input_size=input_size,
+            hidden_sizes=[256, 128, 64, 32],
+            output_size=1,
+            dropout_rate=0.3
+        )
+        model_instance.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+        model_instance.eval()
+        model = model_instance
+
+@predict_bp.route('/predict', methods=['POST'])
+def predict():
+    lazy_load()
+    data = request.get_json()
+    latitude = float(data.get('latitude'))
+    longitude = float(data.get('longitude'))
+    wind_dir = data.get('wind_direction', wind_directions[0]).upper()
+    wind_str = float(data.get('wind_strength', 5))
+
+    # Find nearest beach
+    query_point = np.array([[latitude, longitude]])
+    _, indices = nn_model.kneighbors(query_point)
+    beach_index = indices[0][0]
+    nearest_beach = beach_db.iloc[beach_index]
+
+    # Create input
+    input_data = pd.DataFrame({
+        'Orientation': [nearest_beach['Orientation']],
+        'Sediment': [nearest_beach['Sediment']],
+        'Longitude': [longitude],
+        'Latitude': [latitude],
+        'Wind direction': [wind_dir],
+        'Wind strength': [wind_str]
+    })
+
+    # Preprocess and predict
+    processed_input = preprocessor.transform(input_data)
+    input_tensor = torch.tensor(processed_input, dtype=torch.float32)
+    with torch.no_grad():
+        prediction = model(input_tensor).item()
+
+    # Get nearest beach details for the response (swap lat/lon for frontend map)
+    nearest_beach_details = {
+        'latitude': nearest_beach['Longitude'],   # swap!
+        'longitude': nearest_beach['Latitude'],   # swap!
+        'orientation': nearest_beach['Orientation'],
+        'sediment': nearest_beach['Sediment']
+    }
+
+    return jsonify({
+        'user_latitude': latitude,   # Correct: real latitude
+        'user_longitude': longitude, # Correct: real longitude
+        'wind_direction': wind_dir,
+        'wind_strength': wind_str,
+        'prediction': prediction,
+        'nearest_beach': nearest_beach_details,
+        'success': True
+    })