import streamlit as st
from PIL import Image
import torch
from transformers import AutoModelForImageClassification, AutoFeatureExtractor

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

# Load pre-trained model from Hugging Face
@st.cache_resource
def load_model():
    model_name = "Kaylah072001/stock_prediction_model.h5"  # Replace with your actual model name on Hugging Face
    try:
        model = AutoModelForImageClassification.from_pretrained(model_name)
        feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
        st.success("Model loaded successfully!")
        return model, feature_extractor
    except Exception as e:
        st.error(f"Error loading model: {e}")
        return None, None

model, feature_extractor = load_model()

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

if uploaded_file is not None and model is not None and feature_extractor 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
        inputs = feature_extractor(images=image, return_tensors="pt")

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

        # Display prediction
        st.header("Prediction")
        sentiment = "Bullish" if predicted_class == 1 else "Bearish"
        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.error(f"Error processing image: {e}")