Update app.py

app.py

@@ -1,86 +1,23 @@
-# 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
-from transformers import ViTForImageClassification, ViTFeatureExtractor
 from PIL import Image
-import torch
+import requests
+from transformers import pipeline
 
-# Load the model and feature extractor
-model = ViTForImageClassification.from_pretrained("path/to/your/model")
-feature_extractor = ViTFeatureExtractor.from_pretrained("path/to/your/model")
+# Load the pipeline
+model_name = "vm24bho/net_dfm_myimg"
+pipe = pipeline('image-classification', model=model_name)
 
-st.title("Deepfake Classification App")
+st.title("Deepfake vs Real Image Detection")
 
-# 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.image(image, caption='Uploaded Image.', use_column_width=True)
     st.write("")
     st.write("Classifying...")
 
-    # Preprocess the image
-    inputs = feature_extractor(images=image, return_tensors="pt")
-
-    # Perform inference
-    with torch.no_grad():
-        outputs = model(**inputs)
-
-    # Get the predicted label
-    logits = outputs.logits
-    predicted_class_idx = logits.argmax(-1).item()
-    predicted_class_label = model.config.id2label[predicted_class_idx]
+    # Apply the model
+    result = pipe(image)
 
     # Display the result
-    st.write(f"Prediction: {predicted_class_label}")
+    st.write(result)