import streamlit as st
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.metrics import r2_score
class RegressionModel2(nn.Module):
def __init__(self, input_dim2, hidden_dim2, output_dim2):
super(RegressionModel2, self).__init__()
self.fc1 = nn.Linear(input_dim2, hidden_dim2)
self.relu1 = nn.ReLU()
self.fc2 = nn.Linear(hidden_dim2, output_dim2)
self.batch_norm1 = nn.BatchNorm1d(hidden_dim2)
def forward(self, x2):
out = self.fc1(x2)
out = self.relu1(out)
out = self.batch_norm1(out)
out = self.fc2(out)
return out
# Load the trained model
model2 = RegressionModel2(3, 32, 1)
model2.load_state_dict(torch.load('model.pt'))
model2.eval()
def predict_astigmatism(age, aca_magnitude, aca_axis):
input_data = torch.tensor([[age, aca_magnitude, aca_axis]], dtype=torch.float32)
output = model2(input_data)
return output.item()
def main():
st.set_page_config(page_title='Astigmatism Prediction', page_icon=':eyeglasses:', layout='wide')
st.markdown(
"""
""",
unsafe_allow_html=True
)
st.markdown(
"""
""",
unsafe_allow_html=True
)
age = st.number_input('Enter Patient Age:', min_value=0, step=1)
aca_magnitude = st.number_input('Enter ACA Magnitude:', min_value=0.0, step=0.1)
aca_axis = st.number_input('Enter ACA Axis:', min_value=0, max_value=180, step=1)
if st.button('Predict!'):
astigmatism = predict_astigmatism(age, aca_magnitude, aca_axis)
st.success(f'Predicted Astigmatism: {astigmatism:.4f}')
if __name__ == '__main__':
main()