Update app.py

app.py

@@ -1,46 +1,57 @@
-from transformers import pipeline
 import streamlit as st
-import torch
-import requests
-from PIL import Image
+from transformers import pipeline
 from pytube import YouTube
 import tempfile
+from moviepy.editor import VideoFileClip
+from PIL import Image
+import requests
+from io import BytesIO
+
+st.set_page_config(page_title="AI Media Verifier", layout="wide")
+
+@st.cache_resource
+def load_image_model():
+    return pipeline("image-classification", model="google/vit-base-patch16-224")
+
+@st.cache_resource
+def load_video_model():
+    return pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
 
 def detect_image(image):
-    model = pipeline("image-classification", model="google/vit-base-patch16-224")
+    model = load_image_model()
     return model(image)
 
 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
-
-def main():
-    st.title("VerifiAI - Image & Video Authenticity Checker")
-    option = st.sidebar.selectbox("Select Option", ["Image Detection", "Video Detection"])
-
-    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}%")
-
-    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.")
-
-if __name__ == "__main__":
-    if not torch.cuda.is_available():
-        st.warning("CUDA not available. Running on CPU might be slower.")
-    main()
+    stream = yt.streams.filter(file_extension='mp4', progressive=True).order_by('resolution').desc().first()
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp:
+        stream.download(filename=tmp.name)
+        clip = VideoFileClip(tmp.name).subclip(0, min(10, yt.length))
+        frame = clip.get_frame(clip.duration / 2)
+        image = Image.fromarray(frame)
+        return detect_image(image)
+
+st.title("AI Media Verifier: Image & Video Authenticity")
+
+option = st.radio("Choose input type:", ("Image", "Video"))
+
+if option == "Image":
+    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(uploaded_image)
+        st.success("Analysis Complete")
+        for res in results:
+            st.write(f"**{res['label']}**: {res['score']*100:.2f}%")
+
+elif option == "Video":
+    video_url = st.text_input("Enter YouTube Video Link")
+    if video_url:
+        st.video(video_url)
+        with st.spinner("Analyzing video..."):
+            results = detect_video(video_url)
+        st.success("Video Analysis Complete")
+        for res in results:
+            st.write(f"**{res['label']}**: {res['score']*100:.2f}%")