Spaces:
Running
Running
import gradio as gr | |
import requests | |
import yt_dlp | |
import cv2 | |
from google_img_source_search import ReverseImageSearcher | |
from PIL import Image | |
import os | |
import uuid | |
import subprocess | |
import html | |
import imgbbpy # Adding imgbb to upload images for accessible URLs | |
# Initialize imgbb client with your API key (you need to provide your key) | |
imgbb_client = imgbbpy.SyncClient("YOUR_IMGBB_API_KEY") | |
# Function to download video from a given URL using yt-dlp | |
def dl(inp): | |
out = None | |
try: | |
# Generate a unique ID for each download to avoid conflicts | |
uid = uuid.uuid4() | |
# Replace characters in URL to create a valid filename | |
inp_out = inp.replace("https://", "").replace("/", "_").replace(".", "_").replace("=", "_").replace("?", "_") | |
# Construct yt-dlp command to download the video | |
command = [ | |
'yt-dlp', inp, '--trim-filenames', '160', | |
'-o', f"{uid}/{inp_out}.mp4", '-S', 'res,mp4', '--recode', 'mp4' | |
] | |
# Special handling for Twitter URLs | |
if "twitter" in inp: | |
command.insert(2, '--extractor-arg') | |
command.insert(3, 'twitter:api=syndication') | |
# Run the yt-dlp command to download the video | |
subprocess.run(command, check=True) | |
out = f"{uid}/{inp_out}.mp4" | |
print(out) | |
except subprocess.CalledProcessError as e: | |
print(f"yt-dlp failed: {e}") | |
return out, gr.HTML(""), "", "" | |
# Function to process the video, extracting frames and performing reverse image search | |
def process_vid(file, cur_frame, every_n): | |
uid = uuid.uuid4() # Unique identifier for each run to avoid conflicts | |
capture = cv2.VideoCapture(file) # Open the video file | |
frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT)) # Get total number of frames in the video | |
rev_img_searcher = ReverseImageSearcher() # Initialize the reverse image searcher | |
html_out = "" | |
count = int(every_n) | |
# Determine the starting frame based on user input | |
start_frame = int(cur_frame) if cur_frame else 0 | |
try: | |
# Iterate over frames from the starting frame to the end | |
for i in range(start_frame, frame_count): | |
if count >= int(every_n): | |
count = 1 | |
capture.set(cv2.CAP_PROP_POS_FRAMES, i) # Set the position to the specific frame | |
ret, frame_f = capture.read() # Read the frame | |
if not ret: | |
continue # If frame could not be read, continue to the next one | |
# Save the current frame as an image file | |
frame_path = f"{uid}-vid_tmp{i}.png" | |
cv2.imwrite(frame_path, frame_f) | |
# Upload the frame to imgbb to get a publicly accessible URL | |
response = imgbb_client.upload(file=frame_path) | |
out_url = response.url # Get the public URL from imgbb response | |
# Perform reverse image search on the extracted frame | |
res = rev_img_searcher.search(out_url) | |
out_cnt = 0 | |
# If results are found, generate HTML output | |
if len(res) > 0: | |
for search_item in res: | |
out_cnt += 1 | |
html_out += f""" | |
<div> | |
Title: {html.escape(search_item.page_title)}<br> | |
Site: <a href='{html.escape(search_item.page_url)}' target='_blank' rel='noopener noreferrer'>{html.escape(search_item.page_url)}</a><br> | |
Img: <a href='{html.escape(search_item.image_url)}' target='_blank' rel='noopener noreferrer'>{html.escape(search_item.image_url)}</a><br> | |
<img class='my_im' src='{html.escape(search_item.image_url)}'><br> | |
</div>""" | |
return gr.HTML(f'<h1>Total Found: {out_cnt}</h1><br>{html_out}'), f"Found frame: {i}", i + int(every_n) | |
count += 1 | |
except Exception as e: | |
return gr.HTML(f'{e}'), "", "" | |
return gr.HTML('No frame matches found.'), "", "" | |
# Function to process an image file and convert it for reverse image search | |
def process_im(file, url): | |
# Check if the URL starts with the expected prefix | |
if not url.startswith("https://omnibus"): | |
return url | |
else: | |
# Save the image to a temporary file | |
uid = uuid.uuid4() | |
read_file = Image.open(file) | |
read_file.save(f"{uid}-tmp.png") | |
action_input = f"{uid}-tmp.png" | |
out = os.path.abspath(action_input) | |
# Upload the image to imgbb to get a publicly accessible URL | |
response = imgbb_client.upload(file=out) | |
out_url = response.url # Get the public URL from imgbb response | |
return out_url | |
# Function to perform reverse image search on a given image | |
def rev_im(image): | |
uid = uuid.uuid4() # Generate a unique ID for each run | |
html_out = """""" # Initialize HTML output | |
# Read the image using OpenCV and save it to a temporary file | |
image = cv2.imread(image) | |
tmp_image_path = f"{uid}-im_tmp.png" | |
cv2.imwrite(tmp_image_path, image) | |
# Upload the image to imgbb to get a publicly accessible URL | |
response = imgbb_client.upload(file=tmp_image_path) | |
out_url = response.url # Get the public URL from imgbb response | |
# Perform reverse image search | |
rev_img_searcher = ReverseImageSearcher() | |
res = rev_img_searcher.search(out_url) | |
count = 0 | |
# If results are found, generate HTML output | |
for search_item in res: | |
count += 1 | |
html_out += f""" | |
<div> | |
Title: {html.escape(search_item.page_title)}<br> | |
Site: <a href='{html.escape(search_item.page_url)}' target='_blank' rel='noopener noreferrer'>{html.escape(search_item.page_url)}</a><br> | |
Img: <a href='{html.escape(search_item.image_url)}' target='_blank' rel='noopener noreferrer'>{html.escape(search_item.image_url)}</a><br> | |
<img class='my_im' src='{html.escape(search_item.image_url)}'><br> | |
</div>""" | |
return gr.HTML(f'<h1>Total Found: {count}</h1><br>{html_out}') | |
# Define the Gradio interface for the application | |
with gr.Blocks() as app: | |
with gr.Row(): | |
gr.Column() | |
with gr.Column(): | |
# Radio button to choose between Image or Video input | |
source_tog = gr.Radio(choices=["Image", "Video"], value="Image") | |
# Box to handle image-related input | |
with gr.Box(visible=True) as im_box: | |
inp_url = gr.Textbox(label="Image URL") | |
load_im_btn = gr.Button("Load Image") | |
inp_im = gr.Image(label="Search Image", type='filepath') | |
go_btn_im = gr.Button("Start Image Search") | |
# Box to handle video-related input | |
with gr.Box(visible=False) as vid_box: | |
vid_url = gr.Textbox(label="Video URL") | |
vid_url_btn = gr.Button("Load URL") | |
inp_vid = gr.Video(label="Search Video") | |
with gr.Row(): | |
every_n = gr.Number(label="Every /nth frame", value=10) # Input to select frame extraction frequency | |
stat_box = gr.Textbox(label="Status") | |
with gr.Row(): | |
go_btn_vid = gr.Button("Start Video Search") | |
next_btn = gr.Button("Next Frame Search") | |
gr.Column() | |
with gr.Row(): | |
html_out = gr.HTML(""" | |
""") | |
with gr.Row(visible=False): | |
hid_box = gr.Textbox() | |
# Function to shuffle between image and video input boxes based on user selection | |
def shuf(tog): | |
if tog == "Image": | |
return gr.update(visible=True), gr.update(visible=False) | |
if tog == "Video": | |
return gr.update(visible=False), gr.update(visible=True) | |
# Function to load image from the URL | |
def load_image(url): | |
return url | |
# Gradio button interactions and linking functions | |
im_load = load_im_btn.click(load_image, inp_url, inp_im) | |
next_btn.click(process_vid, [inp_vid, hid_box, every_n], [html_out, stat_box, hid_box]) | |
vid_load = vid_url_btn.click(dl, vid_url, [inp_vid, html_out, stat_box, hid_box]) | |
vid_proc = go_btn_vid.click(process_vid, [inp_vid, hid_box, every_n], [html_out, stat_box, hid_box]) | |
im_proc = go_btn_im.click(rev_im, inp_im, [html_out]) | |
source_tog.change(shuf, [source_tog], [im_box, vid_box], cancels=[vid_proc, im_proc, im_load, vid_load]) | |
# Launch the Gradio app with concurrency settings | |
app.queue(concurrency_count=20).launch() |