Update app.py

app.py

@@ -1,40 +1,56 @@
 import streamlit as st
-import torch
-from model import ToxicImageClassifier
-from utils import load_image, predict_toxicity, get_label
+from transformers import ViTForImageClassification, ViTImageProcessor
+from utils import load_image_vit, predict_toxicity_vit, get_label
 from PIL import Image
+import torch
+import io
 
-def main():
-    st.title("ToxiScan - Toxic Image Classifier")
-    st.write("Upload an image to check if it contains toxic content")
-
-    # Load model
+def classify_image(uploaded_file):
+    # Load pre-trained ViT model and processor
+    model_name = "google/vit-base-patch16-224"
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = ToxicImageClassifier()
+    model = ViTForImageClassification.from_pretrained(model_name)
+    processor = ViTImageProcessor.from_pretrained(model_name)
+
+    # Modify for binary classification (toxic/non-toxic)
     try:
-        model.load_state_dict(torch.load("toxic_classifier.pth", map_location=device))
+        model.classifier = torch.nn.Linear(model.classifier.in_features, 2)
+        model.load_state_dict(torch.load("toxic_classifier.pth", map_location=device), strict=False)
     except FileNotFoundError:
-        st.error("Model weights (toxic_classifier.pth) not found. Please train the model first.")
-        return
+        st.warning("Using pre-trained ImageNet weights. For toxic classification, upload toxic_classifier.pth.")
     model.to(device)
 
+    # Process image and predict
+    inputs = load_image_vit(uploaded_file, processor)
+    prediction, probabilities = predict_toxicity_vit(model, inputs, device)
+    label = get_label(prediction)
+
+    return label, probabilities
+
+def main():
+    st.title("ToxiScan - Toxic Image Classifier")
+    st.write("Upload an image to detect if it contains toxic content using a pre-trained Vision Transformer.")
+
     # File uploader
     uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
 
     if uploaded_file is not None:
-        # Display image
+        # Display uploaded image
         image = Image.open(uploaded_file)
         st.image(image, caption="Uploaded Image", use_column_width=True)
 
         # Process and predict
         with st.spinner("Analyzing..."):
-            image_tensor = load_image(uploaded_file)
-            prediction, probabilities = predict_toxicity(model, image_tensor, device)
-            label = get_label(prediction)
+            label, probabilities = classify_image(uploaded_file)
 
             # Display results
-            st.write(f"Prediction: **{label}**")
-            st.write(f"Confidence: Toxic: {probabilities[1]:.2%}, Non-Toxic: {probabilities[0]:.2%}")
+            st.subheader("Results")
+            st.write(f"**Prediction:** {label}")
+            st.write(f"**Confidence Scores:**")
+            st.write(f"- Toxic: {probabilities[1]:.2%}")
+            st.write(f"- Non-Toxic: {probabilities[0]:.2%}")
+
+            # Bar chart for visualization
             st.bar_chart({"Toxic": probabilities[1], "Non-Toxic": probabilities[0]})
 
 if __name__ == "__main__":