Update app.py

app.py

@@ -5,11 +5,25 @@ 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
-model = resnet50(pretrained=True)
-model.fc = torch.nn.Linear(model.fc.in_features, 2)  # 2 classes: bullish and bearish
-model.load_state_dict(torch.load('/Users/kaylahoffman/Desktop/full_model.h5'))
-model.eval()
+@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/full_model.pt', 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}")
+        st.warning("Using untrained model. Predictions may not be accurate.")
+    model.eval()
+    return model
+
+model = load_model()
 
 # Define image transformation
 transform = transforms.Compose([
@@ -18,36 +32,34 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
-# Streamlit app
-st.title("Stock Trend Predictor: Bullish or Bearish?")
-
 uploaded_file = st.file_uploader("Upload a stock graph image", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
-    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"
-        color = "red"
-    else:
-        sentiment = "Bullish"
-        color = "green"
-
-    confidence = probabilities[predicted_class].item() * 100
-    st.subheader(f"{sentiment}: {confidence:.2f}%")
-    st.progress(confidence / 100, text=f"{sentiment} Confidence")
-
-else:
-    st.info("Please upload a stock graph image to get a prediction.")
+    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"
+            color = "red"
+        else:
+            sentiment = "Bullish"
+            color = "green"
+
+        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