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

size_js="""
    function imgSize(){
        var myImg = document.getElementsByClassName("my_im");
        var realWidth = myImg.naturalWidth;
        var realHeight = myImg.naturalHeight;
        alert("Original width=" + realWidth + ", " + "Original height=" + realHeight);
    }"""

def dl(inp):
    out = None
    out_file = []
    try:
        # Create a unique identifier for each download
        uid = uuid.uuid4()
        os.makedirs(f"{uid}", exist_ok=True)  # Create a directory for storing the downloaded file

        # Generate a unique output filename
        inp_out = inp.replace("https://", "").replace("/", "_").replace(".", "_").replace("=", "_").replace("?", "_")

        # Download using yt_dlp
        command = [
            "yt-dlp",
            inp,
            "--trim-filenames", "160",
            "-o", f"{uid}/{inp_out}.mp4",
            "-S", "res,mp4",
            "--recode", "mp4"
        ]

        if "twitter" in inp:
            command.insert(2, "--extractor-arg")
            command.insert(3, "twitter:api=syndication")

        # Use subprocess to execute the command and handle errors
        result = subprocess.run(command, capture_output=True, text=True)

        # Check if the command was successful
        if result.returncode != 0:
            print("Error occurred:", result.stderr)
            return None, gr.HTML(f"<p>Error: {result.stderr}</p>"), "", ""

        out = f"{uid}/{inp_out}.mp4"
        print(out)

    except Exception as e:
        print(e)
        return None, gr.HTML(f"<p>Error: {str(e)}</p>"), "", ""

    return out, gr.HTML(""), "", ""

def process_vid(file, cur_frame, every_n):
    new_video_in = str(file)
    capture = cv2.VideoCapture(new_video_in)
    frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
    rev_img_searcher = ReverseImageSearcher()
    html_out = ""
    count = int(every_n)
    if cur_frame == "" or cur_frame is None:
        start_frame = 0
    else:
        start_frame = int(cur_frame)
    try:
        for i in range(start_frame, frame_count - 1):
            if count == int(every_n):
                count = 1
                print(i)
                capture.set(cv2.CAP_PROP_POS_FRAMES, i)
                ret, frame_f = capture.read()
                if not ret:
                    continue
                cv2.imwrite(f"{uid}-vid_tmp{i}.png", frame_f)
                out = os.path.abspath(f"{uid}-vid_tmp{i}.png")
                out_url = f'https://nymbo-reverse-image.hf.space/file={out}'
                print(out)
                res = rev_img_searcher.search(out_url)
                out_cnt = 0
                if len(res) > 0:
                    for search_item in res:
                        out_cnt += 1
                        html_out = f"""{html_out}
                        <div>
                        Title: {search_item.page_title}<br>
                        Site: <a href='{search_item.page_url}' target='_blank' rel='noopener noreferrer'>{search_item.page_url}</a><br>
                        Img: <a href='{search_item.image_url}' target='_blank' rel='noopener noreferrer'>{search_item.image_url}</a><br>
                        <img class='my_im' src='{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
            print(i + 1)
    except Exception as e:
        return gr.HTML(f"{e}"), "", ""
    return gr.HTML('No frame matches found.'), "", ""

def process_im(file, url):
    if not url.startswith("https://nymbo"):
        return url
    else:
        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://nymbo-reverse-image.hf.space/file={out}'
        return out_url

def rev_im(image):
    out_list = []
    html_out = """"""
    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://nymbo-reverse-image.hf.space/file={out}'
    rev_img_searcher = ReverseImageSearcher()
    res = rev_img_searcher.search(out_url)
    count = 0
    for search_item in res:
        count += 1
        html_out = f"""{html_out}
        <div>
        Title: {search_item.page_title}<br>
        Site: <a href='{search_item.page_url}' target='_blank' rel='noopener noreferrer'>{search_item.page_url}</a><br>
        Img: <a href='{search_item.image_url}' target='_blank' rel='noopener noreferrer'>{search_item.image_url}</a><br>
        <img class='my_im' src='{search_item.image_url}'><br>
        </div>"""
    return gr.HTML(f'<h1>Total Found: {count}</h1><br>{html_out}')

with gr.Blocks() as app:
    with gr.Row():
        gr.Column()
        with gr.Column():
            source_tog = gr.Radio(choices=["Image", "Video"], value="Image")
            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")
            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)
                stat_box = gr.Textbox(label="Status")
                with gr.Row():
                    go_btn_vid = gr.Button("Start Video Search")
                    next_btn = gr.Button("Next")
    
    with gr.Row():
        html_out = gr.HTML("""")
    with gr.Row(visible=False):
        hid_box = gr.Textbox()

    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)
    
    def load_image(url):
        return url

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

app.queue(concurrency_count=20).launch()