|
import streamlit as st |
|
import requests |
|
from transformers import pipeline |
|
from PIL import Image |
|
from io import BytesIO |
|
import re |
|
|
|
|
|
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() |
|
|