Update app.py

app.py

@@ -0,0 +1,36 @@
+import streamlit as st
+from transformers import ViTFeatureExtractor, ViTForImageClassification
+from PIL import Image
+import torch
+
+# Load the model and feature extractor
+model_name = "google/vit-base-patch16-224"
+model = ViTForImageClassification.from_pretrained(model_name)
+feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
+
+# Streamlit app
+st.title("Image Classifier")
+st.write("Upload an image to classify it into categories.")
+
+# File uploader
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+
+if uploaded_file:
+    # Load and display the image
+    image = Image.open(uploaded_file).convert("RGB")
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+
+    # Preprocess the image
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    
+    # Perform inference
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits
+        predicted_class_idx = logits.argmax(-1).item()
+    
+    # Get classification label
+    label = model.config.id2label[predicted_class_idx]
+    
+    # Display results
+    st.write(f"Prediction: **{label}**")