Update app.py

app.py

@@ -19,40 +19,51 @@ class RegressionModel2(nn.Module):
         return out
 
 # Load the saved model state dictionaries
-model_j0 = RegressionModel2(3, 32, 1)
-model_j0.load_state_dict(torch.load('j0_model.pt'))
-model_j0.eval()
+model_x = RegressionModel2(3, 32, 1)
+model_x.load_state_dict(torch.load('j0_model.pt'))
+model_x.eval()
 
-model_j45 = RegressionModel2(3, 32, 1)
-model_j45.load_state_dict(torch.load('j45_model.pt'))
-model_j45.eval()
+model_y = RegressionModel2(3, 32, 1)
+model_y.load_state_dict(torch.load('j45_model.pt'))
+model_y.eval()
 
 def predict_components(age, axis, aca):
     """
-    This function predicts both J0 and J45 components using the loaded models.
+    This function predicts both x (J0) and y (J45) components using the loaded models.
     """
     data = torch.tensor([[age, axis, aca]], dtype=torch.float32)
     
     with torch.no_grad():
-        j0_pred = model_j0(data).item()
-        j45_pred = model_j45(data).item()
+        x_pred = model_x(data).item()
+        y_pred = model_y(data).item()
     
-    return j0_pred, j45_pred
+    return x_pred, y_pred
 
-def calculate_magnitude_and_axis(j0, j45):
+def calculate_magnitude_and_axis(x, y):
     """
-    Calculate magnitude and axis from J0 and J45 components.
+    Calculate magnitude and axis from x and y components using Excel formulas.
     """
-    magnitude = 2 * math.sqrt(j0**2 + j45**2)
+    magnitude = math.sqrt(x**2 + y**2)
     
-    axis = 0.5 * math.atan2(j45, j0)
-    axis = math.degrees(axis)
+    # Calculate intermediate axis (in radians)
+    intermediate_axis_rad = 0.5 * math.atan2(y, x)
+    
+    # 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
+        else:
+            final_axis = intermediate_axis_deg + 180
+    else:
+        final_axis = intermediate_axis_deg + 90
     
     # Ensure axis is between 0 and 180
-    if axis < 0:
-        axis += 180
+    final_axis = final_axis % 180
     
-    return magnitude, axis
+    return magnitude, final_axis
 
 def main():
     st.set_page_config(page_title='Astigmatism Prediction', page_icon=':eyeglasses:', layout='wide')
@@ -66,10 +77,14 @@ def main():
 
     if st.button('Predict!'):
         if 18 <= age <= 90 and 0 <= aca_magnitude <= 10 and 0 <= aca_axis <= 180:
-            j0_pred, j45_pred = predict_components(age, aca_axis, aca_magnitude)
-            predicted_magnitude, predicted_axis = calculate_magnitude_and_axis(j0_pred, j45_pred)
-            st.success(f'Predicted Total Corneal Astigmatism Magnitude: {predicted_magnitude:.2f} D')
+            x_pred, y_pred = predict_components(age, aca_axis, aca_magnitude)
+            predicted_magnitude, predicted_axis = calculate_magnitude_and_axis(x_pred, y_pred)
+            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}')
         else:
             st.error('Please ensure all inputs are within the specified ranges.')