Update app.py

app.py

@@ -1,40 +1,53 @@
 import streamlit as st
-from transformers import pipeline
 from PIL import Image
+import torch
+from torchvision import transforms
+from torchvision.models import resnet50
 import io
 
-# Load the image classification pipeline
-try:
-    classifier = pipeline(task="image-classification", model="microsoft/resnet-50")
-except Exception as e:
-    st.error(f"Error loading the model: {e}")
-    st.stop()
+# 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('path_to_your_trained_model.pth'))
+model.eval()
 
-# Set up the Streamlit app title
-st.title("Stock Chart Analyzer: Bullish or Bearish? 📈📉")
+# 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]),
+])
 
-# File uploader for image input
-uploaded_file = st.file_uploader("Upload a stock chart image", type=["jpg", "jpeg", "png"])
+# 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:
-    # Display the uploaded image
-    image = Image.open(uploaded_file)
-    st.image(image, caption="Uploaded Stock Chart", use_column_width=True)
-    
-    # Classify the image
-    with st.spinner("Analyzing the trend..."):
-        predictions = classifier(image)
-    
-    # Display the predictions
-    st.header("Trend Analysis")
-    for p in predictions:
-        sentiment = p['label'].split('_')[0].capitalize()
-        confidence = round(p['score'] * 100, 1)
-        st.subheader(f"{sentiment}: {confidence}%")
-        
-        # Add color-coded bars for visual representation
-        color = "green" if sentiment == "Bullish" else "red"
-        st.progress(confidence / 100, text=f"{sentiment} Confidence")
+    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 chart image to analyze the trend.")
+    st.info("Please upload a stock graph image to get a prediction.")