Create app.py

app.py

@@ -0,0 +1,40 @@
+import streamlit as st
+from transformers import ViTImageProcessor, AutoModelForImageClassification
+from PIL import Image
+import requests
+from io import BytesIO
+
+# Load the model and processor
+processor = ViTImageProcessor.from_pretrained('AdamCodd/vit-base-nsfw-detector')
+model = AutoModelForImageClassification.from_pretrained('AdamCodd/vit-base-nsfw-detector')
+
+# Define prediction function
+def predict_image(image):
+    try:
+        # Process the image and make prediction
+        inputs = processor(images=image, return_tensors="pt")
+        outputs = model(**inputs)
+        logits = outputs.logits
+
+        # Get predicted class
+        predicted_class_idx = logits.argmax(-1).item()
+        predicted_label = model.config.id2label[predicted_class_idx]
+
+        return predicted_label
+    except Exception as e:
+        return str(e)
+
+# Streamlit app
+st.title("NSFW Image Classifier")
+
+# Upload image file
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+if uploaded_file is not None:
+    image = Image.open(uploaded_file)
+    st.image(image, caption='Uploaded Image.', use_column_width=True)
+    st.write("")
+    st.write("Classifying...")
+
+    # Predict and display result
+    prediction = predict_image(image)
+    st.write(f"Predicted Class: {prediction}")