Update app.py

app.py

@@ -24,11 +24,11 @@ class RegressionModel2(nn.Module):
 @st.cache_resource
 def load_models():
     model_j0 = RegressionModel2(3, 32, 1)
-    model_j0.load_state_dict(torch.load('j0_model-2.pt'))
+    model_j0.load_state_dict(torch.load('j0_model.pt'))
     model_j0.eval()
 
     model_j45 = RegressionModel2(3, 32, 1)
-    model_j45.load_state_dict(torch.load('j45_model-2.pt'))
+    model_j45.load_state_dict(torch.load('j45_model.pt'))
     model_j45.eval()
     
     return model_j0, model_j45
@@ -56,10 +56,32 @@ def predict_new_j0_j45(age, aca_magnitude, aca_axis_deg):
 def main():
     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, value=58.0, step=0.1)
-    aca_magnitude = st.number_input('Enter ACA Magnitude (0-10 Diopters):', min_value=0.0, max_value=10.0, value=2.3, step=0.01)
-    aca_axis = st.number_input('Enter ACA Axis (0-180 Degrees):', min_value=0.0, max_value=180.0, value=97.7, step=0.1)
+    # Initialize session state for input values if not already present
+    if 'age' not in st.session_state:
+        st.session_state.age = ''
+    if 'aca_magnitude' not in st.session_state:
+        st.session_state.aca_magnitude = ''
+    if 'aca_axis' not in st.session_state:
+        st.session_state.aca_axis = ''
+
+    # Input fields with session state
+    age = st.number_input('Enter Patient Age (18-90 Years):', 
+                          min_value=18.0, max_value=90.0, 
+                          value=st.session_state.age, 
+                          step=0.1, 
+                          key='age')
+    
+    aca_magnitude = st.number_input('Enter ACA Magnitude (0-10 Diopters):', 
+                                    min_value=0.0, max_value=10.0, 
+                                    value=st.session_state.aca_magnitude, 
+                                    step=0.01, 
+                                    key='aca_magnitude')
+    
+    aca_axis = st.number_input('Enter ACA Axis (0-180 Degrees):', 
+                               min_value=0.0, max_value=180.0, 
+                               value=st.session_state.aca_axis, 
+                               step=0.1, 
+                               key='aca_axis')
 
     if st.button('Predict!'):
         if 18 <= age <= 90 and 0 <= aca_magnitude <= 10 and 0 <= aca_axis <= 180:
@@ -75,7 +97,7 @@ def main():
             if predicted_axis < 0:
                 predicted_axis += 180
             
-            # Display results in a green success box
+            # Display results in green success boxes
             st.success(f'Predicted Total Corneal Astigmatism Magnitude: {predicted_magnitude:.2f} D')
             st.success(f'Predicted Total Corneal Astigmatism Axis: {predicted_axis:.1f}°')