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()
 
@@ -34,7 +34,64 @@ def calculate_initial_j0_j45(magnitude, axis):
     j45 = magnitude * math.sin(2 * axis_rad)
     return j0, j45
 
-def predict_new_j0_j45(j0, j45):
+def predict_new_j0_j45(age, magnitude, axis):
+    """Predict new J0 and J45 using the loaded models."""
+    input_data = torch.tensor([[age, magnitude, axis]], dtype=torch.float32)
+    with torch.no_grad():
+        new_j0 = model_j0(input_data).item()
+        new_j45 = model_j45(input_data).item()
+    return new_j0, new_j45
+
+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)
+    
+    if j0 > 0:
+        if j45 > 0:
+            final_axis = z_deg
+        else:
+            final_axis = z_deg + 180
+    else:
+        final_axis = z_deg + 90
+    
+    final_axis = final_axis % 180
+    
+    return magnitude, final_axis
+
+def main():
+    st.set_page_config(page_title='Astigmatism Prediction', page_icon=':eyeglasses:', layout='wide')
+    
+    st.title('Total Corneal Astigmatism Prediction')
+
+    # Input fields
+    age = st.number_input('Enter Patient Age (18-90 Years):', min_value=18.0, max_value=90.0, step=0.1)
+    aca_magnitude = st.number_input('Enter ACA Magnitude (0-10 Diopters):', min_value=0.0, max_value=10.0, step=0.01)
+    aca_axis = st.number_input('Enter ACA Axis (0-180 Degrees):', min_value=0.0, max_value=180.0, step=0.1)
+
+    if st.button('Predict!'):
+        if 18 <= age <= 90 and 0 <= aca_magnitude <= 10 and 0 <= aca_axis <= 180:
+            # Step 2: Calculate initial J0 and J45 (for display purposes only)
+            initial_j0, initial_j45 = calculate_initial_j0_j45(aca_magnitude, aca_axis)
+            
+            # Step 3: Predict new J0 and J45 using the models
+            new_j0, new_j45 = predict_new_j0_j45(age, aca_magnitude, aca_axis)
+            
+            # 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 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.')
+
+if __name__ == '__main__':
+    main()def predict_new_j0_j45(j0, j45):
     """Predict new J0 and J45 using the loaded models."""
     with torch.no_grad():
         new_j0 = model_j0(torch.tensor([[j0]], dtype=torch.float32)).item()