Update app.py

app.py

@@ -1,63 +1,49 @@
 import streamlit as st
 from PIL import Image
 import torch
-from torchvision import transforms
-from torchvision.models import resnet50
-import io
+from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 
 # Streamlit app
 st.title("Stock Trend Predictor: Bullish or Bearish?")
 
-# Load pre-trained ResNet50 model
+# Load pre-trained model from Hugging Face
 @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
+    model_name = "Kaylah072001/stock_prediction_model"  # Replace with your actual model name on Hugging Face
     try:
-        state_dict = torch.load('Kaylah072001/stock_prediction_model.h5', map_location=torch.device('cpu'))
-        model.load_state_dict(state_dict)
+        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}")
-    model.eval()
-    return model
+        return None, None
 
-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]),
-])
+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:
+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
-        input_tensor = transform(image).unsqueeze(0)
+        inputs = feature_extractor(images=image, return_tensors="pt")
 
         # Make prediction
         with torch.no_grad():
-            output = model(input_tensor)
-            probabilities = torch.nn.functional.softmax(output[0], dim=0)
+            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")
-        if predicted_class == 0:
-            sentiment = "Bearish"
-
-        else:
-            sentiment = "Bullish"
-
+        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
+        st.error(f"Error processing image: {e}")