Update app.py

app.py

@@ -1,16 +1,18 @@
 import streamlit as st
 from PIL import Image
-from transformers import pipeline
+import torch
+from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 
-# Load the image classification pipeline
-model_name = "vikramjeetthakur/hotornot"  # Replace with your model's name or path
-classifier = pipeline("image-classification", model=model_name)
+# Load a pre-trained model and feature extractor
+model_name = "facebook/wide_resnet50_2"  # Using a general model
+model = AutoModelForImageClassification.from_pretrained(model_name)
+feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
 
 # Define the main function for the Streamlit app
 def main():
     st.title("Hot or Not Image Classifier")
 
-    st.write("Upload an image to classify it as hot or not hot.")
+    st.write("Upload an image to classify it.")
 
     # Image upload
     uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
@@ -20,14 +22,18 @@ def main():
         image = Image.open(uploaded_file)
         st.image(image, caption="Uploaded Image", use_column_width=True)
 
-        # Make predictions
-        predictions = classifier(image)
+        # Preprocess the image
+        inputs = feature_extractor(images=image, return_tensors="pt")
 
-        # Display results
-        for prediction in predictions:
-            label = prediction['label']
-            score = prediction['score']
-            st.write(f"**{label}** with confidence {score:.2f}")
+        # Make predictions
+        with torch.no_grad():
+            outputs = model(**inputs)
+            logits = outputs.logits  # Get the logits
+            class_idx = logits.argmax(-1).item()  # Get the index of the highest probability
+
+        # Display results based on class index
+        st.write(f"Predicted class index: {class_idx}")
+        st.write(f"Predicted class label: {model.config.id2label[class_idx]}")
 
 if __name__ == "__main__":
     main()