Update app.py

app.py

@@ -1,50 +1,49 @@
 import streamlit as st
-from yt_dlp import YoutubeDL
+import requests
 from transformers import pipeline
 from PIL import Image
-import requests
 from io import BytesIO
+import re
 
-@st.cache_resource
-def load_model():
-    return pipeline("zero-shot-image-classification", model="openai/clip-vit-base-patch16")
+# Initialize image classification pipeline
+model = pipeline("image-classification", model="google/vit-base-patch16-224")
 
-def fetch_thumbnail(video_url):
-    ydl_opts = {
-        'skip_download': True,
-        'quiet': True,
-        'extract_flat': True,
-        'cookiefile': 'cookies.txt'
-    }
-    with YoutubeDL(ydl_opts) as ydl:
-        info = ydl.extract_info(video_url, download=False)
-        return info.get('title'), info.get('thumbnail')
+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 classify_thumbnail(model, thumbnail_url):
+def detect_thumbnail(thumbnail_url):
+    """Run AI detection on the thumbnail image."""
     response = requests.get(thumbnail_url)
-    image = Image.open(BytesIO(response.content)).convert("RGB")
-    labels = ["real", "AI-generated", "manipulated", "deepfake"]
-    return model(image, candidate_labels=labels)
+    if response.status_code == 200:
+        image = Image.open(BytesIO(response.content))
+        results = model(image)
+        return results, image
+    return None, None
 
 def main():
-    st.title("🎥 Video Integrity Checker")
-    video_url = st.text_input("Enter YouTube video URL:")
-    if st.button("Submit") and video_url:
-        try:
-            title, thumbnail = fetch_thumbnail(video_url)
-            st.success(f"Video Title: {title}")
-            if thumbnail:
-                st.image(thumbnail, caption="Video Thumbnail", use_container_width=True)
-                model = load_model()
-                with st.spinner("Analyzing..."):
-                    results = classify_thumbnail(model, thumbnail)
+    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 label, score in zip(results['labels'], results['scores']):
-                    st.write(f"{label}: {score:.2%}")
+                for res in results:
+                    st.write(f"- **{res['label']}**: {res['score']:.4f}")
             else:
-                st.warning("No thumbnail found.")
-        except Exception as e:
-            st.error(f"Error: {e}")
+                st.error("Failed to retrieve or analyze the thumbnail.")
+        else:
+            st.error("Invalid YouTube URL.")
 
 if __name__ == "__main__":
     main()