Update app.py

app.py

@@ -18,15 +18,14 @@ class RegressionModel2(nn.Module):
         out = self.fc2(out)
         return out
 
-# Load the saved model state dictionaries
 @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
@@ -34,14 +33,12 @@ def load_models():
 model_j0, model_j45 = load_models()
 
 def calculate_initial_j0_j45(magnitude, axis_deg):
-    """Calculate initial J0 and J45 from magnitude and axis (in degrees)."""
     axis_rad = math.radians(axis_deg)
     j0 = magnitude * math.cos(2 * axis_rad)
     j45 = magnitude * math.sin(2 * axis_rad)
     return j0, j45
 
 def predict_new_j0_j45(age, aca_magnitude, aca_axis_deg):
-    """Predict new J0 and J45 using the loaded models."""
     initial_j0, initial_j45 = calculate_initial_j0_j45(aca_magnitude, aca_axis_deg)
     
     input_data_j0 = torch.tensor([[age, aca_axis_deg, initial_j0]], dtype=torch.float32)
@@ -54,45 +51,51 @@ def predict_new_j0_j45(age, aca_magnitude, aca_axis_deg):
     return new_j0, new_j45
 
 def main():
-    st.title('Total Corneal Astigmatism Prediction')
+    st.set_page_config(page_title='Total Corneal Astigmatism Prediction', page_icon=':eyeglasses:', layout='wide')
     
-    # User input fields
-    age = st.number_input('Enter Patient Age (18-90 Years):', min_value=18.0, max_value=90.0, value=58.0, step=1.0)
-    aca_magnitude = st.number_input('Enter ACA Magnitude (0-10 Diopters):', min_value=0.0, max_value=10.0, value=2.3, step=0.1)
+    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)
-    
-    if st.button('Predict'):
-        # Calculate initial J0 and J45
-        initial_j0, initial_j45 = calculate_initial_j0_j45(aca_magnitude, aca_axis)
-        
-        # Make prediction
-        new_j0, new_j45 = predict_new_j0_j45(age, aca_magnitude, aca_axis)
-        
-        # Calculate TCA magnitude and axis
-        tca_magnitude = math.sqrt(new_j0**2 + new_j45**2)
-        tca_axis = 0.5 * math.degrees(math.atan2(new_j45, new_j0))
-        if tca_axis < 0:
-            tca_axis += 180
-        
-        # Display results
-        st.subheader('Prediction Results')
-        col1, col2 = st.columns(2)
-        with col1:
-            st.write(f"Initial J0: {initial_j0:.2f}")
-            st.write(f"Initial J45: {initial_j45:.2f}")
-            st.write(f"Predicted J0: {new_j0:.2f}")
-            st.write(f"Predicted J45: {new_j45:.2f}")
-        with col2:
-            st.write(f"Predicted TCA Magnitude: {tca_magnitude:.2f} D")
-            st.write(f"Predicted TCA Axis: {tca_axis:.1f}°")
 
-        # Optional: Display input tensors and raw outputs for verification
-        if st.checkbox('Show detailed model inputs and outputs'):
-            st.subheader('Model Details')
-            st.write("Input tensor for J0:", torch.tensor([[age, aca_axis, initial_j0]], dtype=torch.float32))
-            st.write("Input tensor for J45:", torch.tensor([[age, aca_axis, initial_j45]], dtype=torch.float32))
-            st.write("Raw J0 output:", new_j0)
-            st.write("Raw J45 output:", new_j45)
+    if st.button('Predict!'):
+        if 18 <= age <= 90 and 0 <= aca_magnitude <= 10 and 0 <= aca_axis <= 180:
+            # Calculate initial J0 and J45
+            initial_j0, initial_j45 = calculate_initial_j0_j45(aca_magnitude, aca_axis)
+            
+            # Predict new J0 and J45 using the models
+            new_j0, new_j45 = predict_new_j0_j45(age, aca_magnitude, aca_axis)
+            
+            # Calculate predicted magnitude and axis
+            predicted_magnitude = math.sqrt(new_j0**2 + new_j45**2)
+            predicted_axis = 0.5 * math.degrees(math.atan2(new_j45, new_j0))
+            if predicted_axis < 0:
+                predicted_axis += 180
+            
+            # Display results in a green success box
+            st.success(f'''
+            Predicted Total Corneal Astigmatism Magnitude: {predicted_magnitude:.2f} D
+            Predicted Total Corneal Astigmatism Axis: {predicted_axis:.1f}°
+            ''')
+            
+            # Display intermediate values for verification
+            st.info(f'''
+            Input ACA - Magnitude: {aca_magnitude:.2f} D, Axis: {aca_axis:.1f}°
+            Initial J0: {initial_j0:.2f}, Initial J45: {initial_j45:.2f}
+            Predicted J0: {new_j0:.2f}, Predicted J45: {new_j45:.2f}
+            ''')
+            
+            # Additional debugging information (optional)
+            if st.checkbox('Show detailed model inputs and outputs'):
+                st.subheader('Debugging Information:')
+                st.write(f"Input tensor for J0: {torch.tensor([[age, aca_axis, initial_j0]], dtype=torch.float32)}")
+                st.write(f"Input tensor for J45: {torch.tensor([[age, aca_axis, initial_j45]], dtype=torch.float32)}")
+                st.write(f"Raw J0 output: {new_j0}")
+                st.write(f"Raw J45 output: {new_j45}")
+        else:
+            st.error('Please ensure all inputs are within the specified ranges.')
 
 if __name__ == '__main__':
     main()