app.py

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()