Update app.py

app.py

@@ -1,41 +1,46 @@
 import streamlit as st
 from PIL import Image
-from transformers import pipeline
 import torch
-import tensorflow as tf
+import torchvision.transforms as transforms
+from torchvision import models
 
-# Check if PyTorch or TensorFlow is available
-if not (torch.cuda.is_available() or torch.__version__):
-    st.warning("PyTorch is not installed or not available.")
-if not (tf.__version__):
-    st.warning("TensorFlow is not installed.")
+# Check if CUDA (GPU) is available, otherwise use CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load the pre-trained model (ResNet-50)
+model = models.resnet50(pretrained=True)
+model = model.to(device)
+model.eval()
+
+# Define a function to preprocess the image
+def preprocess_image(image):
+    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]),
+    ])
+    return transform(image).unsqueeze(0).to(device)
 
 # Set up a title for the app
 st.title("Image Recognition App")
 
-# Load a pre-trained image classification pipeline from Hugging Face
-try:
-    classifier = pipeline('image-classification', model='google/efficientnet-b0')
-except RuntimeError as e:
-    st.error("Error loading the model: Please ensure PyTorch or TensorFlow is installed.")
-    st.error(str(e))
-
 # Upload an image
 uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
 
 if uploaded_file is not None:
     # Display the uploaded image
-    image = Image.open(uploaded_file)
+    image = Image.open(uploaded_file).convert("RGB")
     st.image(image, caption='Uploaded Image.', use_column_width=True)
-    
-    # Run the classifier
+
+    # Preprocess the image and make predictions
     st.write("Classifying...")
-    try:
-        results = classifier(image)
-        # Display the classification results
+    input_tensor = preprocess_image(image)
+    with torch.no_grad():
+        output = model(input_tensor)
+        probabilities = torch.nn.functional.softmax(output[0], dim=0)
+        top5_prob, top5_classes = torch.topk(probabilities, 5)
+        
+        # Display the top 5 predictions
         st.write("Predictions:")
-        for label in results:
-            st.write(f"Label: {label['label']}, Confidence: {label['score']:.2f}")
-    except Exception as e:
-        st.error("An error occurred during classification.")
-        st.error(str(e))
+        for i in range(5):
+            st.write(f"Label: {top5_classes[i].item()}, Confidence: {top5_prob[i].item():.2f}")