VidGuard / app.py
TruthLens's picture
Update app.py
67ddf1e verified
raw
history blame
1.77 kB
import streamlit as st
import requests
from transformers import pipeline
from PIL import Image
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:
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.")
else:
st.error("Invalid YouTube URL.")
if __name__ == "__main__":
main()