Spaces:
Runtime error
Runtime error
File size: 4,225 Bytes
8c89017 1e2e161 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 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
}) |