Update app.py

app.py

@@ -1,64 +1,56 @@
 import streamlit as st
-from transformers import pipeline
-from PIL import Image
 import requests
+from PIL import Image
+from transformers import pipeline
 from urllib.parse import urlparse, parse_qs
-from io import BytesIO
 
-# Initialize the deepfake detection model
+st.set_page_config(page_title="Deepfake Video Detector", layout="centered")
+
+st.title("🎥 Deepfake Video Detector")
+st.write("Enter a YouTube video URL (supports Shorts & standard videos) to check for deepfakes.")
+
+@st.cache_data
+def extract_video_id(url):
+    parsed_url = urlparse(url)
+    if "youtube" in parsed_url.netloc:
+        if parsed_url.path.startswith("/watch"):
+            return parse_qs(parsed_url.query).get("v", [None])[0]
+        elif parsed_url.path.startswith("/shorts/"):
+            return parsed_url.path.split("/")[-1]
+    return None
+
+@st.cache_data
+def get_thumbnail(video_id):
+    return f"https://img.youtube.com/vi/{video_id}/hqdefault.jpg"
+
 @st.cache_resource
 def load_model():
-    return pipeline("image-classification", model="Wvolf/ViT_Deepfake_Detection")
+    return pipeline("image-classification", model="facebook/deit-base-distilled-patch16-224")
 
 model = load_model()
 
-def get_thumbnail_url(video_url):
-    """
-    Extracts the YouTube video ID and returns the thumbnail URL.
-    """
-    parsed_url = urlparse(video_url)
-    video_id = None
-    if 'youtube' in parsed_url.netloc:
-        query_params = parse_qs(parsed_url.query)
-        video_id = query_params.get('v', [None])[0]
-    elif 'youtu.be' in parsed_url.netloc:
-        video_id = parsed_url.path.lstrip('/')
+video_url = st.text_input("🔗 Paste YouTube video URL:")
+if st.button("Detect Deepfake") and video_url:
+    video_id = extract_video_id(video_url)
     
     if video_id:
-        return f"https://img.youtube.com/vi/{video_id}/maxresdefault.jpg"
-    return None
+        thumbnail_url = get_thumbnail(video_id)
+        try:
+            response = requests.get(thumbnail_url)
+            response.raise_for_status()
+            thumbnail = Image.open(requests.get(thumbnail_url, stream=True).raw)
+            st.image(thumbnail, caption="Video Thumbnail", use_container_width=True)
 
-def analyze_thumbnail(thumbnail_url):
-    """
-    Downloads the thumbnail image and analyzes it using the deepfake detection model.
-    """
-    response = requests.get(thumbnail_url)
-    if response.status_code == 200:
-        image = Image.open(BytesIO(response.content)).convert("RGB")
-        results = model(image)
-        return results, image
-    else:
-        st.error("Failed to retrieve the thumbnail image.")
-        return None, None
-
-# Streamlit UI
-st.title("Deepfake Detection from YouTube Thumbnails")
-video_url = st.text_input("Enter YouTube Video URL:")
-if st.button("Analyze"):
-    if video_url:
-        thumbnail_url = get_thumbnail_url(video_url)
-        if thumbnail_url:
-            results, image = analyze_thumbnail(thumbnail_url)
-            if results and image:
-                st.image(image, caption="YouTube Video Thumbnail", use_column_width=True)
-                st.subheader("Detection Results:")
-                for result in results:
-                    label = result['label']
-                    confidence = result['score'] * 100
-                    st.write(f"**{label}**: {confidence:.2f}%")
+            st.write("🔎 Analyzing thumbnail...")
+            results = model(thumbnail)
+            
+            deepfake_score = sum(result['score'] for result in results if 'fake' in result['label'].lower())
+            if deepfake_score > 0.5:
+                st.error(f"⚠️ High probability of deepfake detected! (Confidence: {deepfake_score:.2%})")
             else:
-                st.error("Could not analyze the thumbnail.")
-        else:
-            st.error("Invalid YouTube URL. Please enter a valid URL.")
+                st.success(f"✅ No strong evidence of deepfake detected. (Confidence: {1 - deepfake_score:.2%})")
+
+        except requests.exceptions.RequestException:
+            st.error("❌ Failed to fetch thumbnail. Please check the video URL.")
     else:
-        st.warning("Please enter a YouTube video URL.")
+        st.error("❌ Invalid YouTube URL. Please enter a valid video or Shorts link.")