Update app.py

app.py

@@ -1,47 +1,86 @@
+# import streamlit as st
+# import torch
+# from torchvision.transforms import transforms
+# from PIL import Image
+# from transformers import AutoModelForSequenceClassification
+# # Load the model and tokenizer
+# model_name = "vm24bho/net_dfm_myimg"
+# model = AutoModelForSequenceClassification.from_pretrained(model_name)
+
+
+# # Define transformations for the input 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]),
+# ])
+
+# def predict(image):
+#     # Preprocess the image
+#     image = transform(image).unsqueeze(0)  # Add batch dimension
+
+#     # Perform inference
+#     outputs = model(image)
+
+#     # Get prediction
+#     prediction = torch.argmax(outputs.logits).item()
+    
+#     return prediction
+
+# def main():
+#     st.title("Image Detection: Real or Deepfake")
+#     uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+
+#     if uploaded_image is not None:
+#         image = Image.open(uploaded_image)
+#         st.image(image, caption='Uploaded Image', use_column_width=True)
+
+#         # Make prediction
+#         if st.button("Detect"):
+#             prediction = predict(image)
+#             if prediction == 0:
+#                 st.write("Prediction: Real")
+#             else:
+#                 st.write("Prediction: Deepfake")
+
+# if __name__ == "__main__":
+#     main()
+
+
 import streamlit as st
-import torch
-from torchvision.transforms import transforms
+from transformers import ViTForImageClassification, ViTFeatureExtractor
 from PIL import Image
-from transformers import AutoModelForSequenceClassification
-# Load the model and tokenizer
-model_name = "vm24bho/net_dfm_myimg"
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
+import torch
+
+# Load the model and feature extractor
+model = ViTForImageClassification.from_pretrained("path/to/your/model")
+feature_extractor = ViTFeatureExtractor.from_pretrained("path/to/your/model")
 
+st.title("Deepfake Classification App")
 
-# Define transformations for the input 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]),
-])
+# File uploader
+uploaded_file = st.file_uploader("Choose an image...", type="jpg")
+
+if uploaded_file is not None:
+    # Load the image
+    image = Image.open(uploaded_file)
+
+    # Display the image
+    st.image(image, caption='Uploaded Image', use_column_width=True)
+    st.write("")
+    st.write("Classifying...")
 
-def predict(image):
     # Preprocess the image
-    image = transform(image).unsqueeze(0)  # Add batch dimension
+    inputs = feature_extractor(images=image, return_tensors="pt")
 
     # Perform inference
-    outputs = model(image)
+    with torch.no_grad():
+        outputs = model(**inputs)
 
-    # Get prediction
-    prediction = torch.argmax(outputs.logits).item()
-    
-    return prediction
-
-def main():
-    st.title("Image Detection: Real or Deepfake")
-    uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-
-    if uploaded_image is not None:
-        image = Image.open(uploaded_image)
-        st.image(image, caption='Uploaded Image', use_column_width=True)
-
-        # Make prediction
-        if st.button("Detect"):
-            prediction = predict(image)
-            if prediction == 0:
-                st.write("Prediction: Real")
-            else:
-                st.write("Prediction: Deepfake")
-
-if __name__ == "__main__":
-    main()
+    # Get the predicted label
+    logits = outputs.logits
+    predicted_class_idx = logits.argmax(-1).item()
+    predicted_class_label = model.config.id2label[predicted_class_idx]
+
+    # Display the result
+    st.write(f"Prediction: {predicted_class_label}")