Kaylah072001 commited on
Commit
71a5a36
·
verified ·
1 Parent(s): a7acc43

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +65 -0
app.py ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from PIL import Image
3
+ import torch
4
+ from torchvision import transforms
5
+ from torchvision.models import resnet50
6
+ import io
7
+
8
+ # Streamlit app
9
+ st.title("Stock Trend Predictor: Bullish or Bearish?")
10
+
11
+ # Load pre-trained ResNet50 model
12
+ @st.cache_resource
13
+ def load_model():
14
+ model = resnet50(pretrained=True)
15
+ model.fc = torch.nn.Linear(model.fc.in_features, 2) # 2 classes: bullish and bearish
16
+ try:
17
+ state_dict = torch.load('/Users/kaylahoffman/Desktop/full_model.pt', map_location=torch.device('cpu'))
18
+ model.load_state_dict(state_dict)
19
+ st.success("Model loaded successfully!")
20
+ except Exception as e:
21
+ st.error(f"Error loading model: {e}")
22
+ st.warning("Using untrained model. Predictions may not be accurate.")
23
+ model.eval()
24
+ return model
25
+
26
+ model = load_model()
27
+
28
+ # Define image transformation
29
+ transform = transforms.Compose([
30
+ transforms.Resize((224, 224)),
31
+ transforms.ToTensor(),
32
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
33
+ ])
34
+
35
+ uploaded_file = st.file_uploader("Upload a stock graph image", type=["jpg", "jpeg", "png"])
36
+
37
+ if uploaded_file is not None:
38
+ try:
39
+ image = Image.open(uploaded_file).convert('RGB')
40
+ st.image(image, caption="Uploaded Stock Graph", use_column_width=True)
41
+
42
+ # Preprocess the image
43
+ input_tensor = transform(image).unsqueeze(0)
44
+
45
+ # Make prediction
46
+ with torch.no_grad():
47
+ output = model(input_tensor)
48
+ probabilities = torch.nn.functional.softmax(output[0], dim=0)
49
+ predicted_class = torch.argmax(probabilities).item()
50
+
51
+ # Display prediction
52
+ st.header("Prediction")
53
+ if predicted_class == 0:
54
+ sentiment = "Bearish"
55
+ color = "red"
56
+ else:
57
+ sentiment = "Bullish"
58
+ color = "green"
59
+
60
+ confidence = probabilities[predicted_class].item() * 100
61
+ st.subheader(f"{sentiment}: {confidence:.2f}%")
62
+ st.progress(confidence / 100, text=f"{sentiment} Confidence")
63
+
64
+ except Exception as e:
65
+ st