Update app.py

app.py

@@ -19,11 +19,11 @@ class RegressionModel2(nn.Module):
         return out
 
 # Load the saved model state dictionaries
-model_j0 = RegressionModel2(1, 32, 1)
+model_j0 = RegressionModel2(3, 32, 1)
 model_j0.load_state_dict(torch.load('j0_model.pt'))
 model_j0.eval()
 
-model_j45 = RegressionModel2(1, 32, 1)
+model_j45 = RegressionModel2(3, 32, 1)
 model_j45.load_state_dict(torch.load('j45_model.pt'))
 model_j45.eval()
 
@@ -36,17 +36,16 @@ def calculate_initial_j0_j45(magnitude, axis_deg):
 
 def predict_new_j0_j45(age, aca_magnitude, aca_axis_deg):
     """Predict new J0 and J45 using the loaded models."""
-    # Calculate initial J0 and J45
-    initial_j0, initial_j45 = calculate_initial_j0_j45(aca_magnitude, aca_axis_deg)
+    aca_axis_rad = math.radians(aca_axis_deg)
+    aca_x = aca_magnitude * math.cos(aca_axis_rad)
+    aca_y = aca_magnitude * math.sin(aca_axis_rad)
     
-    # Prepare inputs for each model
-    input_data_j0 = torch.tensor([[initial_j0]], dtype=torch.float32)
-    input_data_j45 = torch.tensor([[initial_j45]], dtype=torch.float32)
+    input_data_j0 = torch.tensor([[age, aca_axis_deg, aca_x]], dtype=torch.float32)
+    input_data_j45 = torch.tensor([[age, aca_axis_deg, aca_y]], dtype=torch.float32)
     
     with torch.no_grad():
         new_j0 = model_j0(input_data_j0).item()
         new_j45 = model_j45(input_data_j45).item()
-    
     return new_j0, new_j45
 
 def calculate_magnitude_and_axis(j0, j45):