Update app.py

app.py

@@ -1,49 +1,63 @@
 import streamlit as st
 from PIL import Image
 import torch
-from transformers import AutoModelForImageClassification, AutoFeatureExtractor
+from torchvision import transforms
+from torchvision.models import resnet50
+import io
 
 # Streamlit app
 st.title("Stock Trend Predictor: Bullish or Bearish?")
 
-# Load pre-trained model from Hugging Face
+# Load pre-trained ResNet50 model
 @st.cache_resource
 def load_model():
-    model_name = "Kaylah072001/stock_prediction_model.h5"  # Replace with your actual model name on Hugging Face
+    model = resnet50(pretrained=True)
+    model.fc = torch.nn.Linear(model.fc.in_features, 2)  # 2 classes: bullish and bearish
     try:
-        model = AutoModelForImageClassification.from_pretrained(model_name)
-        feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+        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!")
-        return model, feature_extractor
     except Exception as e:
         st.error(f"Error loading model: {e}")
-        return None, None
+    model.eval()
+    return model
 
-model, feature_extractor = load_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 and model is not None and feature_extractor is not None:
+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
-        inputs = feature_extractor(images=image, return_tensors="pt")
+        input_tensor = transform(image).unsqueeze(0)
 
         # Make prediction
         with torch.no_grad():
-            outputs = model(**inputs)
-            logits = outputs.logits
-            probabilities = torch.nn.functional.softmax(logits[0], dim=0)
+            output = model(input_tensor)
+            probabilities = torch.nn.functional.softmax(output[0], dim=0)
             predicted_class = torch.argmax(probabilities).item()
 
         # Display prediction
         st.header("Prediction")
-        sentiment = "Bullish" if predicted_class == 1 else "Bearish"
+        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.error(f"Error processing image: {e}")
+        st