import streamlit as st
from PIL import Image
import torch
from torchvision import transforms
from torchvision.models import resnet50
import io

# Streamlit app
st.title("Stock Trend Predictor: Bullish or Bearish?")

# Load pre-trained ResNet50 model
@st.cache_resource
def load_model():
    model = resnet50(pretrained=True)
    model.fc = torch.nn.Linear(model.fc.in_features, 2)  # 2 classes: bullish and bearish
    try:
        state_dict = torch.load('/Users/kaylahoffman/Desktop/stock_prediction_model.h5', map_location=torch.device('cpu'))
        model.load_state_dict(state_dict)
        st.success("Model loaded successfully!")
    except Exception as e:
        st.error(f"Error loading model: {e}")
    model.eval()
    return model

model = load_model()

# Define image transformation
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

uploaded_file = st.file_uploader("Upload a stock graph image", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    try:
        image = Image.open(uploaded_file).convert('RGB')
        st.image(image, caption="Uploaded Stock Graph", use_column_width=True)

        # Preprocess the image
        input_tensor = transform(image).unsqueeze(0)

        # Make prediction
        with torch.no_grad():
            output = model(input_tensor)
            probabilities = torch.nn.functional.softmax(output[0], dim=0)
            predicted_class = torch.argmax(probabilities).item()

        # Display prediction
        st.header("Prediction")
        if predicted_class == 0:
            sentiment = "Bearish"

        else:
            sentiment = "Bullish"

        confidence = probabilities[predicted_class].item() * 100
        st.subheader(f"{sentiment}: {confidence:.2f}%")
        st.progress(confidence / 100, text=f"{sentiment} Confidence")

    except Exception as e:
        st