Update app.py

app.py

@@ -1,49 +1,46 @@
 import streamlit as st
-import requests
-from transformers import pipeline
 from PIL import Image
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+import torch
+import requests
 from io import BytesIO
-import re
-
-# Initialize image classification pipeline
-model = pipeline("image-classification", model="google/vit-base-patch16-224")
-
-def get_thumbnail_url(youtube_url):
-    """Extract YouTube video ID and generate thumbnail URL."""
-    video_id_match = re.search(r"(?:v=|youtu.be/|shorts/)([a-zA-Z0-9_-]{11})", youtube_url)
-    if video_id_match:
-        video_id = video_id_match.group(1)
-        return f"https://img.youtube.com/vi/{video_id}/maxresdefault.jpg"
-    return None
-
-def detect_thumbnail(thumbnail_url):
-    """Run AI detection on the thumbnail image."""
-    response = requests.get(thumbnail_url)
-    if response.status_code == 200:
+
+# 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)
+
+st.title("Deepfake Thumbnail Detector")
+st.write("Upload a YouTube video link, and we’ll analyze the thumbnail to check for deepfakes.")
+
+video_url = st.text_input("Enter YouTube Video Link:")
+submit = st.button("Detect Deepfake")
+
+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()
+
         image = Image.open(BytesIO(response.content))
-        results = model(image)
-        return results, image
-    return None, None
-
-def main():
-    st.title("🔍 Video Thumbnail AI Detector")
-    st.write("Enter a YouTube video link to detect content from its thumbnail.")
-
-    video_url = st.text_input("YouTube Video URL:")
-    
-    if st.button("Analyze Thumbnail") and video_url:
-        thumbnail_url = get_thumbnail_url(video_url)
-        if thumbnail_url:
-            results, image = detect_thumbnail(thumbnail_url)
-            if results and image:
-                st.image(image, caption="Video Thumbnail", use_container_width=True)
-                st.subheader("Detection Results:")
-                for res in results:
-                    st.write(f"- **{res['label']}**: {res['score']:.4f}")
-            else:
-                st.error("Failed to retrieve or analyze the thumbnail.")
+        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.")
         else:
-            st.error("Invalid YouTube URL.")
+            st.success(f"✅ **No Deepfake Detected** (Confidence: {confidence:.2%})")
+            st.write("👍 Thumbnail appears authentic based on the model’s analysis.")
 
-if __name__ == "__main__":
-    main()
+    except Exception as e:
+        st.error(f"Error: {str(e)}")