Update app.py

app.py

@@ -1,4 +1,3 @@
-# Final version...
 import torch
 import torch.nn as nn
 import gradio as gr
@@ -28,23 +27,25 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 def add_time_decimal_feature(df):
     """
     Add 'time_decimal' feature by combining 'hour' and 'minutes'.
-
     :param df: DataFrame with 'hour' and 'minutes' columns.
     :return: DataFrame with 'time_decimal' and without 'hour' and 'minutes'.
     """
-    if 'hour' in df.columns and 'minutes' in df.columns:
+    if 'time_decimal' in df.columns:
+        logging.info("'time_decimal' feature already exists. Skipping creation.")
+        return df
+    elif 'hour' in df.columns and 'minutes' in df.columns:
         logging.info("Adding 'time_decimal' feature...")
         df['time_decimal'] = df['hour'] + df['minutes'] / 60.0
         df = df.drop(columns=['hour', 'minutes'])  # Drop 'hour' and 'minutes' after creation
         logging.info("'time_decimal' feature added.")
     else:
-        logging.warning("'hour' and/or 'minutes' columns not found. Skipping 'time_decimal' feature addition.")
+        logging.warning("Neither 'time_decimal' nor 'hour' and 'minutes' columns found. Cannot create 'time_decimal' feature.")
+        raise ValueError("Input data must contain 'time_decimal' or both 'hour' and 'minutes' columns.")
     return df
 
 def haversine(lon1, lat1, lon2, lat2):
     """
     Calculate the great-circle distance between two points on the Earth.
-
     :param lon1: Longitude of point 1 (in decimal degrees)
     :param lat1: Latitude of point 1 (in decimal degrees)
     :param lon2: Longitude of point 2 (in decimal degrees)
@@ -65,7 +66,6 @@ def haversine(lon1, lat1, lon2, lat2):
 def calculate_bearing(lon1, lat1, lon2, lat2):
     """
     Calculate the bearing between two points.
-
     :param lon1: Longitude of point 1 (in decimal degrees)
     :param lat1: Latitude of point 1 (in decimal degrees)
     :param lon2: Longitude of point 2 (in decimal degrees)
@@ -90,7 +90,6 @@ def calculate_bearing(lon1, lat1, lon2, lat2):
 def angular_divergence(bearing1, bearing2):
     """
     Calculate the smallest angle difference between two bearings.
-
     :param bearing1: First bearing in degrees
     :param bearing2: Second bearing in degrees
     :return: Angular divergence in degrees
@@ -101,7 +100,6 @@ def angular_divergence(bearing1, bearing2):
 def denormalize(scaled_lat, scaled_lon, scaler, lat_idx, lon_idx):
     """
     Denormalize latitude and longitude using the scaler's parameters.
-
     :param scaled_lat: Scaled latitude values (numpy array).
     :param scaled_lon: Scaled longitude values (numpy array).
     :param scaler: The scaler object used for normalization.
@@ -118,7 +116,7 @@ def denormalize(scaled_lat, scaled_lon, scaler, lat_idx, lon_idx):
     denorm_lon = scaled_lon * (lon_max - lon_min) + lon_min
     return denorm_lat, denorm_lon
 
-def create_dataset_grouped_by_mmsi(df_scaled, seq_len, forecast_horizon, features_to_scale):
+def create_dataset_grouped_by_mmsi(df_scaled, seq_len, forecast_horizon, features_to_scale, future_features):
     """
     Create input and output sequences grouped by original MMSI.
     Returns scaled last known positions.
@@ -137,14 +135,14 @@ def create_dataset_grouped_by_mmsi(df_scaled, seq_len, forecast_horizon, feature
                 # Future positions to predict (scaled)
                 future_positions = group[['latitude_degrees', 'longitude_degrees']].iloc[i + seq_len:i + seq_len + forecast_horizon].to_numpy()
 
-                # Future hour feature
-                future_hour = group[['time_decimal']].iloc[i + seq_len].values[0]
-                future_hour_feature = np.full((seq_len, 1), future_hour)
+                # Future features
+                future_feature_values = group[future_features].iloc[i + seq_len].values
+                future_feature_array = np.tile(future_feature_values, (seq_len, 1))
 
-                # Combine sequence with future_hour_feature
-                sequence_with_future_hour = np.hstack((sequence, future_hour_feature))
+                # Combine sequence with future features
+                sequence_with_future_features = np.hstack((sequence, future_feature_array))
 
-                Xs.append(sequence_with_future_hour)
+                Xs.append(sequence_with_future_features)
                 ys.append(future_positions)
                 mmsis.append(mmsi)
 
@@ -163,7 +161,6 @@ class LSTMModelTeacher(nn.Module):
     def __init__(self, in_dim, hidden_dim, forecast_horizon, n_layers=7, dropout=0.2):
         """
         Teacher LSTM Model.
-
         :param in_dim: Number of input features.
         :param hidden_dim: Number of hidden units.
         :param forecast_horizon: Number of future steps to predict.
@@ -187,7 +184,6 @@ class LSTMModelStudent(nn.Module):
     def __init__(self, in_dim, hidden_dim, forecast_horizon, n_layers=3, dropout=0.2):
         """
         Student LSTM Model.
-
         :param in_dim: Number of input features.
         :param hidden_dim: Number of hidden units.
         :param forecast_horizon: Number of future steps to predict.
@@ -213,48 +209,58 @@ class LSTMModelStudent(nn.Module):
 
 def load_models(model_paths):
     """
-    Load teacher and student models, including submodels for North, Mid, and South areas.
-
+    Load teacher, student, and cargo vessel models, including submodels for North, Mid, and South areas.
     :param model_paths: Dictionary containing paths to the models.
     :return: Dictionary of loaded models.
     """
     models = {}
     logging.info("Loading Teacher model...")
+    # Teacher model input dimension
+    teacher_in_dim = 15  # Features including 'future_hour_feature' (time_decimal)
     # Load Teacher Model (Global)
-    teacher = LSTMModelTeacher(in_dim=15, hidden_dim=200, forecast_horizon=1, n_layers=7, dropout=0.2)  # 15 features including 'future_hour_feature'
+    teacher = LSTMModelTeacher(in_dim=teacher_in_dim, hidden_dim=200, forecast_horizon=1, n_layers=7, dropout=0.2)
     teacher.load_state_dict(torch.load(model_paths['teacher'], map_location=torch.device('cpu')))
     teacher.eval()
     models['Teacher'] = teacher
     logging.info("Teacher model loaded successfully.")
 
     logging.info("Loading Student North model...")
-    # Load Student Models (Sub-areas)
-    student_north = LSTMModelStudent(in_dim=15, hidden_dim=200, forecast_horizon=1, n_layers=3, dropout=0.2)
+    # Student North model input dimension is the same as teacher
+    student_north = LSTMModelStudent(in_dim=teacher_in_dim, hidden_dim=200, forecast_horizon=1, n_layers=3, dropout=0.2)
     student_north.load_state_dict(torch.load(model_paths['student_north'], map_location=torch.device('cpu')))
     student_north.eval()
     models['Student_North'] = student_north
     logging.info("Student North model loaded successfully.")
 
     logging.info("Loading Student Mid model...")
-    student_mid = LSTMModelStudent(in_dim=15, hidden_dim=200, forecast_horizon=1, n_layers=3, dropout=0.2)
+    student_mid = LSTMModelStudent(in_dim=teacher_in_dim, hidden_dim=200, forecast_horizon=1, n_layers=3, dropout=0.2)
     student_mid.load_state_dict(torch.load(model_paths['student_mid'], map_location=torch.device('cpu')))
     student_mid.eval()
     models['Student_Mid'] = student_mid
     logging.info("Student Mid model loaded successfully.")
 
     logging.info("Loading Student South model...")
-    student_south = LSTMModelStudent(in_dim=15, hidden_dim=200, forecast_horizon=1, n_layers=3, dropout=0.2)
+    student_south = LSTMModelStudent(in_dim=teacher_in_dim, hidden_dim=200, forecast_horizon=1, n_layers=3, dropout=0.2)
     student_south.load_state_dict(torch.load(model_paths['student_south'], map_location=torch.device('cpu')))
     student_south.eval()
     models['Student_South'] = student_south
     logging.info("Student South model loaded successfully.")
 
+    # Load Cargo Vessel model
+    logging.info("Loading Cargo Vessel model...")
+    # Cargo Vessel model input dimension
+    cargo_in_dim = 13 + 3  # 13 features (without 'year') + 3 future features ('day', 'month', 'time_decimal')
+    cargo_model = LSTMModelTeacher(in_dim=cargo_in_dim, hidden_dim=200, forecast_horizon=1, n_layers=10, dropout=0.2)
+    cargo_model.load_state_dict(torch.load(model_paths['cargo_vessel'], map_location=torch.device('cpu')))
+    cargo_model.eval()
+    models['Cargo_Vessel'] = cargo_model
+    logging.info("Cargo Vessel model loaded successfully.")
+
     return models
 
 def load_scalers(scaler_paths):
     """
     Load scalers for each model.
-
     :param scaler_paths: Dictionary containing paths to the scaler files.
     :return: Dictionary of loaded scalers.
     """
@@ -275,7 +281,6 @@ def load_scalers(scaler_paths):
 def determine_subarea(df):
     """
     Determine the sub-area (North, Mid, South) based on latitude and longitude ranges.
-
     :param df: DataFrame containing 'latitude_degrees' and 'longitude_degrees'.
     :return: String indicating the sub-area.
     """
@@ -307,24 +312,27 @@ def determine_subarea(df):
 
     return predominant_subarea
 
-def select_model(models, subarea):
+def select_model(models, subarea, model_choice):
     """
-    Select the appropriate model based on the sub-area.
-
+    Select the appropriate model based on the sub-area and model choice.
     :param models: Dictionary of loaded models.
     :param subarea: String indicating the sub-area.
+    :param model_choice: String indicating the selected model.
     :return: Tuple of (selected_model, selected_model_name).
     """
-    if subarea in ['North', 'Mid', 'South']:
-        selected_model = models.get(f'Student_{subarea}')
-        selected_model_name = f'Student_{subarea}'
-        logging.info(f"Selected model: {selected_model_name}")
-        return selected_model, selected_model_name
+    if model_choice == "Auto-Select":
+        if subarea in ['North', 'Mid', 'South']:
+            selected_model = models.get(f'Student_{subarea}')
+            selected_model_name = f'Student_{subarea}'
+        else:
+            selected_model = models.get('Teacher')
+            selected_model_name = 'Teacher'
     else:
-        selected_model = models.get('Teacher')
-        selected_model_name = 'Teacher'
-        logging.info(f"Selected model: {selected_model_name}")
-        return selected_model, selected_model_name
+        selected_model = models.get(model_choice)
+        selected_model_name = model_choice
+
+    logging.info(f"Selected model: {selected_model_name}")
+    return selected_model, selected_model_name
 
 # ============================
 # Evaluation Metrics Calculation
@@ -333,7 +341,6 @@ def select_model(models, subarea):
 def calculate_classic_metrics(y_true, y_pred):
     """
     Calculate MAE, MSE, and RMSE directly on latitude/longitude pairs.
-
     :param y_true: Ground truth positions (numpy array of shape (num_samples, 2)).
     :param y_pred: Predicted positions (numpy array of shape (num_samples, 2)).
     :return: Dictionary containing the classic metrics.
@@ -360,7 +367,6 @@ def calculate_classic_metrics(y_true, y_pred):
 def calculate_distance_metrics(y_true, y_pred):
     """
     Calculate metrics based on distance (in kilometers).
-
     :param y_true: Ground truth positions (numpy array of shape (num_samples, 2)).
     :param y_pred: Predicted positions (numpy array of shape (num_samples, 2)).
     :return: Dictionary containing the distance-based metrics.
@@ -416,35 +422,13 @@ def classical_prediction(file_path, model_choice, min_mmsi, max_mmsi, models, lo
         if df.empty:
             error_message = "No data available after applying MMSI filters."
             logging.error(error_message)
-            return {"error": error_message}, None, None
-
-        # Check if 'time_decimal' exists
-        if 'time_decimal' not in df.columns:
-            df = add_time_decimal_feature(df)
-        else:
-            logging.info("'time_decimal' feature already exists. Skipping creation.")
-
-        expected_columns = [
-            "mmsi", "sog_kt", "latitude_degrees", "longitude_degrees", "cog_degrees",
-            "dimension_a_m", "dimension_b_m", "dimension_c_m", "dimension_d_m",
-            "ship_type", "day", "month", "year", "time_decimal"
-        ]
-        if list(df.columns) != expected_columns:
-            error_message = (
-                f"Input data does not have the correct columns.\n"
-                f"Expected columns: {expected_columns}\n"
-                f"Got columns: {list(df.columns)}"
-            )
-            logging.error(error_message)
-            return {"error": error_message}, None, None
-
-        logging.info("Input CSV has the correct columns.")
+            return {"error": error_message}, None, None, None
 
         # Select the appropriate model and scaler
         if model_choice == "Auto-Select":
             temp_df = df.copy()
             subarea = determine_subarea(temp_df)
-            selected_model, selected_model_name = select_model(models, subarea)
+            selected_model, selected_model_name = select_model(models, subarea, model_choice)
             scaler = loaded_scalers[selected_model_name]
         else:
             if model_choice in models:
@@ -454,17 +438,50 @@ def classical_prediction(file_path, model_choice, min_mmsi, max_mmsi, models, lo
             else:
                 error_message = f"Selected model '{model_choice}' is not available."
                 logging.error(error_message)
-                return {"error": error_message}, None, None
+                return {"error": error_message}, None, None, None
 
         logging.info(f"Using scaler for model: {selected_model_name}")
 
+        # Adjust features_to_scale based on the selected model
+        if selected_model_name == 'Cargo_Vessel':
+            features_to_scale = [
+                "mmsi", "sog_kt", "latitude_degrees", "longitude_degrees", "cog_degrees",
+                "dimension_a_m", "dimension_b_m", "dimension_c_m", "dimension_d_m",
+                "ship_type", "day", "month", "time_decimal"
+            ]
+            future_features = ['day', 'month', 'time_decimal']
+        else:
+            features_to_scale = [
+                "mmsi", "sog_kt", "latitude_degrees", "longitude_degrees", "cog_degrees",
+                "dimension_a_m", "dimension_b_m", "dimension_c_m", "dimension_d_m",
+                "ship_type", "day", "month", "year", "time_decimal"
+            ]
+            future_features = ['time_decimal']
+
+        # Check if the necessary columns exist
+        expected_columns = features_to_scale
+        if not all(col in df.columns for col in expected_columns):
+            error_message = (
+                f"Input data does not have the correct columns.\n"
+                f"Expected columns for {selected_model_name}: {expected_columns}\n"
+                f"Got columns: {list(df.columns)}"
+            )
+            logging.error(error_message)
+            return {"error": error_message}, None, None, None
+
+        logging.info("Input CSV has the correct columns.")
+
+        # Check and add 'time_decimal' if necessary
+        if selected_model_name != 'Cargo_Vessel':
+            df = add_time_decimal_feature(df)
+        else:
+            if 'time_decimal' not in df.columns:
+                error_message = "Cargo model requires 'time_decimal' column."
+                logging.error(error_message)
+                return {"error": error_message}, None, None, None
+
         # Normalize the data
         logging.info("Normalizing the data...")
-        features_to_scale = [
-            "mmsi", "sog_kt", "latitude_degrees", "longitude_degrees", "cog_degrees",
-            "dimension_a_m", "dimension_b_m", "dimension_c_m", "dimension_d_m",
-            "ship_type", "day", "month", "year", "time_decimal"
-        ]
         X_new = df[features_to_scale]
         X_scaled = scaler.transform(X_new)
         df_scaled = pd.DataFrame(X_scaled, columns=features_to_scale, index=df.index)
@@ -473,12 +490,14 @@ def classical_prediction(file_path, model_choice, min_mmsi, max_mmsi, models, lo
         # Create sequences and get last known positions (scaled)
         seq_len = 24
         forecast_horizon = 1
-        X, y, mmsi_seq, last_known_positions_scaled = create_dataset_grouped_by_mmsi(df_scaled, seq_len, forecast_horizon, features_to_scale)
+        X, y, mmsi_seq, last_known_positions_scaled = create_dataset_grouped_by_mmsi(
+            df_scaled, seq_len, forecast_horizon, features_to_scale, future_features
+        )
 
         if X.size == 0:
             error_message = "Not enough data to create sequences."
             logging.error(error_message)
-            return {"error": error_message}, None, None
+            return {"error": error_message}, None, None, None
 
         logging.info(f"Created {X.shape[0]} sequences.")
 
@@ -572,9 +591,13 @@ def abnormal_behavior_detection(prediction_file_path, alpha=0.5, threshold=10.0)
 
         # Check if necessary columns exist
         expected_columns = [
-            "mmsi", "sog_kt", "latitude_degrees", "longitude_degrees", "cog_degrees",
-            "dimension_a_m", "dimension_b_m", "dimension_c_m", "dimension_d_m",
-            "ship_type", "day", "month", "year", "time_decimal"
+            'MMSI',
+            'Last Known Latitude',
+            'Last Known Longitude',
+            'Predicted Latitude',
+            'Predicted Longitude',
+            'Real Latitude',
+            'Real Longitude'
         ]
 
         if not all(col in df.columns for col in expected_columns):
@@ -584,7 +607,7 @@ def abnormal_behavior_detection(prediction_file_path, alpha=0.5, threshold=10.0)
                 f"Got columns: {list(df.columns)}"
             )
             logging.error(error_message)
-            return {"error": error_message}
+            return None, error_message
 
         # Extract necessary data
         mmsi_seq = df['MMSI'].values
@@ -660,15 +683,14 @@ def abnormal_behavior_detection(prediction_file_path, alpha=0.5, threshold=10.0)
 # ============================
 
 def main():
-    # ============================
-    # Define Model and Scaler Paths
-    # ============================
+
 
     model_paths = {
-        'teacher': 'LSTM_whole_atlantic_horizon1_with_time_decimal_input_batch256/horizon_data_LSTM_whole_atlantic_horizon1_with_time_decimal_input_batch256_seq_24/run_1/best_model.pth',
+         'teacher': 'LSTM_whole_atlantic_horizon1_with_time_decimal_input_batch256/horizon_data_LSTM_whole_atlantic_horizon1_with_time_decimal_input_batch256_seq_24/run_1/best_model.pth',
         'student_north': 'LSTM_whole_atlantic_horizon1_with_time_decimal_input_batch256_KD_North/horizon1_data_LSTM_whole_atlantic_horizon1_with_time_decimal_input_batch256_KD_North_seq_24/run_1/best_model.pth',
         'student_mid': 'LSTM_whole_atlantic_horizon1_with_time_decimal_input_batch256_KD_Mid/horizon1_data_LSTM_whole_atlantic_horizon1_with_time_decimal_input_batch256_KD_Mid_seq_24/run_1/best_model.pth',
-        'student_south': 'LSTM_whole_atlantic_horizon1_with_time_decimal_input_batch256_KD_South/horizon1_data_LSTM_whole_atlantic_horizon1_with_time_decimal_input_batch256_KD_South_seq_24/run_1/best_model.pth'
+        'student_south': 'LSTM_whole_atlantic_horizon1_with_time_decimal_input_batch256_KD_South/horizon1_data_LSTM_whole_atlantic_horizon1_with_time_decimal_input_batch256_KD_South_seq_24/run_1/best_model.pth',
+        'cargo_vessel': 'Cago_final/LSTMModel_cargo_horizon1_with_month_day_time_input_batch256/horizon_data_LSTMModel_cargo_horizon1_with_month_day_time_input_batch256_seq_24/run_1/best_model.pth' 
     }
 
     scaler_paths = {
@@ -676,11 +698,10 @@ def main():
         'Student_North': 'scaler_train_North_up.joblib',
         'Student_Mid': 'scaler_train_Mid_up.joblib',
         'Student_South': 'scaler_train_South_up.joblib'
+        'Cargo_Vessel': 'scaler_features_cargo_up_final.joblib'  # Add this line
     }
 
-    # ============================
-    # Load Models and Scalers
-    # ============================
+-
 
     logging.info("Loading models and scalers...")
     models = load_models(model_paths)
@@ -693,7 +714,7 @@ def main():
         inputs=[
             gr.File(label="Upload CSV File", type='filepath'),
             gr.Dropdown(
-                choices=["Auto-Select", "Teacher", "Student_North", "Student_Mid", "Student_South"],
+                choices=["Auto-Select", "Teacher", "Student_North", "Student_Mid", "Student_South", "Cargo_Vessel"],
                 value="Auto-Select",
                 label="Choose Model"
             ),
@@ -749,5 +770,4 @@ def main():
 
 # Run the app
 if __name__ == "__main__":
-    main()  
-    
+    main()