Update app.py

app.py

@@ -19,48 +19,42 @@ class RegressionModel2(nn.Module):
         return out
 
 # Load the saved model state dictionaries
-model_x = RegressionModel2(3, 32, 1)
-model_x.load_state_dict(torch.load('j0_model.pt'))
-model_x.eval()
+model_j0 = RegressionModel2(1, 32, 1)
+model_j0.load_state_dict(torch.load('j0_model.pt'))
+model_j0.eval()
 
-model_y = RegressionModel2(3, 32, 1)
-model_y.load_state_dict(torch.load('j45_model.pt'))
-model_y.eval()
+model_j45 = RegressionModel2(1, 32, 1)
+model_j45.load_state_dict(torch.load('j45_model.pt'))
+model_j45.eval()
 
-def predict_components(age, axis, aca):
-    """
-    This function predicts both x (J0) and y (J45) components using the loaded models.
-    """
-    data = torch.tensor([[age, axis, aca]], dtype=torch.float32)
-    
+def calculate_initial_j0_j45(magnitude, axis):
+    """Calculate initial J0 and J45 from magnitude and axis."""
+    axis_rad = math.radians(axis)
+    j0 = magnitude * math.cos(2 * axis_rad)
+    j45 = magnitude * math.sin(2 * axis_rad)
+    return j0, j45
+
+def predict_new_j0_j45(j0, j45):
+    """Predict new J0 and J45 using the loaded models."""
     with torch.no_grad():
-        x_pred = model_x(data).item()
-        y_pred = model_y(data).item()
-    
-    return x_pred, y_pred
+        new_j0 = model_j0(torch.tensor([[j0]], dtype=torch.float32)).item()
+        new_j45 = model_j45(torch.tensor([[j45]], dtype=torch.float32)).item()
+    return new_j0, new_j45
 
-def calculate_magnitude_and_axis(x, y):
-    """
-    Calculate magnitude and axis from x and y components using Excel formulas.
-    """
-    magnitude = math.sqrt(x**2 + y**2)
-    
-    # Calculate intermediate axis (in radians)
-    intermediate_axis_rad = 0.5 * math.atan2(y, x)
+def calculate_magnitude_and_axis(j0, j45):
+    """Calculate magnitude and axis from J0 and J45."""
+    magnitude = math.sqrt(j0**2 + j45**2)
+    z = 0.5 * math.atan2(j45, j0)
+    z_deg = math.degrees(z)
     
-    # Convert to degrees
-    intermediate_axis_deg = math.degrees(intermediate_axis_rad)
-    
-    # Calculate final axis
-    if x > 0:
-        if y > 0:
-            final_axis = intermediate_axis_deg
+    if j0 > 0:
+        if j45 > 0:
+            final_axis = z_deg
         else:
-            final_axis = intermediate_axis_deg + 180
+            final_axis = z_deg + 180
     else:
-        final_axis = intermediate_axis_deg + 90
+        final_axis = z_deg + 90
     
-    # Ensure axis is between 0 and 180
     final_axis = final_axis % 180
     
     return magnitude, final_axis
@@ -77,14 +71,21 @@ def main():
 
     if st.button('Predict!'):
         if 18 <= age <= 90 and 0 <= aca_magnitude <= 10 and 0 <= aca_axis <= 180:
-            x_pred, y_pred = predict_components(age, aca_axis, aca_magnitude)
-            predicted_magnitude, predicted_axis = calculate_magnitude_and_axis(x_pred, y_pred)
+            # Step 2: Calculate initial J0 and J45
+            initial_j0, initial_j45 = calculate_initial_j0_j45(aca_magnitude, aca_axis)
+            
+            # Step 3: Predict new J0 and J45
+            new_j0, new_j45 = predict_new_j0_j45(initial_j0, initial_j45)
+            
+            # Steps 4-6: Calculate predicted magnitude and axis
+            predicted_magnitude, predicted_axis = calculate_magnitude_and_axis(new_j0, new_j45)
+            
             st.success(f'Predicted Total Corneal Astigmatism Magnitude: {predicted_magnitude:.4f} D')
             st.success(f'Predicted Total Corneal Astigmatism Axis: {predicted_axis:.1f}°')
             
-            # Display predicted x and y components for verification
-            st.info(f'Predicted x (J0) component: {x_pred:.4f}')
-            st.info(f'Predicted y (J45) component: {y_pred:.4f}')
+            # Display intermediate values for verification
+            st.info(f'Initial J0: {initial_j0:.4f}, Initial J45: {initial_j45:.4f}')
+            st.info(f'Predicted J0: {new_j0:.4f}, Predicted J45: {new_j45:.4f}')
         else:
             st.error('Please ensure all inputs are within the specified ranges.')