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
# 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)
out_url = f'https://omnibus-reverse-image.hf.space/file={frame_path}'
# 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"""
"""
return gr.HTML(f'Total Found: {out_cnt}
{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)
out_url = f'https://omnibus-reverse-image.hf.space/file={out}'
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)
cv2.imwrite(f"{uid}-im_tmp.png", image)
out = os.path.abspath(f"{uid}-im_tmp.png")
out_url = f'https://omnibus-reverse-image.hf.space/file={out}'
# 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"""
"""
return gr.HTML(f'Total Found: {count}
{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()