Update app.py

app.py

@@ -1,46 +1,64 @@
 import streamlit as st
+from transformers import pipeline
 from PIL import Image
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification
-import torch
 import requests
+from urllib.parse import urlparse, parse_qs
 from io import BytesIO
 
-# Load deepfake detection model and extractor
-model_name = "systemkc/deepfake-detection-v1"  # Example model (replace if needed)
-extractor = AutoFeatureExtractor.from_pretrained(model_name)
-model = AutoModelForImageClassification.from_pretrained(model_name)
+# Initialize the deepfake detection model
+@st.cache_resource
+def load_model():
+    return pipeline("image-classification", model="Wvolf/ViT_Deepfake_Detection")
 
-st.title("Deepfake Thumbnail Detector")
-st.write("Upload a YouTube video link, and we’ll analyze the thumbnail to check for deepfakes.")
+model = load_model()
 
-video_url = st.text_input("Enter YouTube Video Link:")
-submit = st.button("Detect Deepfake")
+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('/')
+    
+    if video_id:
+        return f"https://img.youtube.com/vi/{video_id}/maxresdefault.jpg"
+    return None
 
-if submit and video_url:
-    try:
-        video_id = video_url.split("v=")[-1].split("&")[0] if "v=" in video_url else video_url.split("/")[-1]
-        thumbnail_url = f"https://img.youtube.com/vi/{video_id}/maxresdefault.jpg"
-        response = requests.get(thumbnail_url)
-        response.raise_for_status()
+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
 
-        image = Image.open(BytesIO(response.content))
-        st.image(image, caption="Video Thumbnail", use_container_width=True)
-
-        # Preprocess and predict
-        inputs = extractor(images=image, return_tensors="pt")
-        with torch.no_grad():
-            outputs = model(**inputs)
-            logits = outputs.logits
-            predicted_class = torch.argmax(logits, dim=1).item()
-            confidence = torch.softmax(logits, dim=1)[0, predicted_class].item()
-
-        # Interpret results
-        if predicted_class == 1:  # Assuming class 1 = Deepfake
-            st.error(f"🚨 **Deepfake Detected** (Confidence: {confidence:.2%})")
-            st.write("⚠️ Indicators of manipulation detected in the thumbnail.")
+# 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}%")
+            else:
+                st.error("Could not analyze the thumbnail.")
         else:
-            st.success(f"✅ **No Deepfake Detected** (Confidence: {confidence:.2%})")
-            st.write("👍 Thumbnail appears authentic based on the model’s analysis.")
-
-    except Exception as e:
-        st.error(f"Error: {str(e)}")
+            st.error("Invalid YouTube URL. Please enter a valid URL.")
+    else:
+        st.warning("Please enter a YouTube video URL.")