Update app.py

app.py

@@ -1,36 +1,46 @@
-import streamlit as st
 from transformers import pipeline
-from PIL import Image
+import streamlit as st
+import torch
 import requests
-from io import BytesIO
+from PIL import Image
+from pytube import YouTube
+import tempfile
+
+def detect_image(image):
+    model = pipeline("image-classification", model="google/vit-base-patch16-224")
+    return model(image)
 
-# Title and description
-st.title("Deepfake Image & Video Detector")
-st.write("Upload an image or enter a video link to check for AI manipulations.")
+def detect_video(video_url):
+    yt = YouTube(video_url)
+    stream = yt.streams.filter(file_extension='mp4').first()
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp_file:
+        stream.download(filename=tmp_file.name)
+        video_path = tmp_file.name
+    return video_path
 
-# Image Detection Function
-def detect_image(image_file):
-    model = pipeline("image-classification", model="microsoft/resnet-50")
-    image = Image.open(image_file)
-    results = model(image)
-    return results
+def main():
+    st.title("VerifiAI - Image & Video Authenticity Checker")
+    option = st.sidebar.selectbox("Select Option", ["Image Detection", "Video Detection"])
 
-# Video Detection Function (Placeholder)
-def detect_video(video_link):
-    st.write("Video detection coming soon.")
+    if option == "Image Detection":
+        uploaded_image = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
+        if uploaded_image:
+            image = Image.open(uploaded_image)
+            st.image(image, caption="Uploaded Image", use_container_width=True)
+            with st.spinner("Analyzing Image..."):
+                results = detect_image(image)
+            for result in results:
+                st.write(f"{result['label']}: {result['score']*100:.2f}%")
 
-# Image Upload Section
-uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
-if uploaded_image:
-    st.image(uploaded_image, caption="Uploaded Image", use_container_width=True)
-    with st.spinner("Analyzing image..."):
-        image_results = detect_image(uploaded_image)
-    st.write("### Detection Results:")
-    for result in image_results:
-        st.write(f"{result['label']}: {result['score']*100:.2f}%")
+    elif option == "Video Detection":
+        video_url = st.text_input("Enter YouTube Video Link")
+        if video_url:
+            st.video(video_url)
+            with st.spinner("Analyzing Video..."):
+                video_path = detect_video(video_url)
+                st.success("Video downloaded for analysis. Video detection model coming soon.")
 
-# Video Link Section
-video_link = st.text_input("Enter a video link")
-if video_link:
-    st.write(f"Analyzing video at: {video_link}")
-    detect_video(video_link)
+if __name__ == "__main__":
+    if not torch.cuda.is_available():
+        st.warning("CUDA not available. Running on CPU might be slower.")
+    main()