import streamlit as st
from utils import *
import torch
import pickle
from PIL import Image

resnetmodel = custom_resnet()
resnetmodel.load_state_dict(torch.load('/app/MIN_RESNET101_BMI_Cache_test.pkl', map_location=torch.device('cpu')))
resnetmodel = resnetmodel.to(device)
resnetmodel.eval()
gpr = pickle.load(open('/app/gpr_model_withgender.pkl', 'rb'))
obj = Data_Processor()


def get_features(img):
    values = []
    image = Image.open(img).convert('RGB')
    values.append(1)
    body_feature = obj.test(image)
    values.append(body_feature.WSR)
    values.append(body_feature.WTR)
    values.append(body_feature.WHpR)
    values.append(body_feature.WHdR)
    values.append(body_feature.HpHdR)
    values.append(body_feature.Area)
    values.append(body_feature.H2W)
    image = Image.open(img).convert('RGB')
    image = ScaleAndPadTransform(224).transform(image)
    image = image.unsqueeze(0)
    data = image.to("cpu")
    conv_out = LayerActivations(resnetmodel.fc1, 1)
    out = resnetmodel(data)
    conv_out.remove()
    xs = torch.squeeze(conv_out.features.cpu().detach()).numpy()

    for x in xs:
        values.append(float(x))

    return values
def main():
    st.title("BMI Prediction App")

    image = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
    if image is not None:
        cols = st.columns(2)  # Create two columns
        with cols[0]:  # Place image in the first column
            st.image(image, caption="Uploaded Image", use_column_width=True)
        
        # Convert image to features
        values = get_features(image)
        
        # Predict BMI using Gaussian Process Regression
        bmi_pred = gpr.predict([values])
        
        with cols[1]:  # Place prediction in the second column
            st.write("Predicted BMI:", bmi_pred[0])
            st.success("Prediction Completed")
        
        st.balloons()
        st.write("<script>window.scrollTo(0, document.body.scrollHeight);</script>", unsafe_allow_html=True)

if __name__ == "__main__":
    main()