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

# Generate a unique identifier for temporary file storage
uid = str(uuid.uuid4())

# Ensure the temporary directory exists
if not os.path.exists(uid):
    os.makedirs(uid)

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);
}
"""

# Function to download video using yt-dlp
def dl(inp):
    out = None
    inp_out = inp.replace("https://", "").replace("/", "_").replace(".", "_").replace("=", "_").replace("?", "_")
    try:
        if "twitter" in inp:
            os.system(f'yt-dlp "{inp}" --extractor-arg "twitter:api=syndication" --trim-filenames 160 -o "{uid}/{inp_out}.mp4" -S res,mp4 --recode mp4')
        else:
            os.system(f'yt-dlp "{inp}" --trim-filenames 160 -o "{uid}/{inp_out}.mp4" -S res,mp4 --recode mp4')

        out = f"{uid}/{inp_out}.mp4"
        if not os.path.exists(out):
            print("Error: Video download failed. File not found.")
            return None, gr.HTML("<h1>Error: Video download failed.</h1>"), "", ""

        print(f"Downloaded video file: {out}")
    except Exception as e:
        print(f"Exception during video download: {e}")
        return None, gr.HTML(f"<h1>Error: {e}</h1>"), "", ""
    return out, gr.HTML(""), "", ""

# Function to process video and perform reverse image search
def process_vid(file, cur_frame, every_n):
    new_video_in = str(file)
    if not os.path.exists(new_video_in):
        print("Error: Video file does not exist.")
        return gr.HTML("<h1>Error: Video file not found.</h1>"), "", ""

    capture = cv2.VideoCapture(new_video_in)
    if not capture.isOpened():
        print("Error: Video file could not be opened.")
        return gr.HTML("<h1>Error: Failed to open video file.</h1>"), "", ""

    frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
    rev_img_searcher = ReverseImageSearcher()
    html_out = ""
    count = int(every_n)

    start_frame = int(cur_frame) if cur_frame else 0

    try:
        for i in range(start_frame, frame_count - 1):
            if count == int(every_n):
                count = 1
                capture.set(cv2.CAP_PROP_POS_FRAMES, i)
                ret, frame_f = capture.read()

                if not ret:
                    print(f"Error: Failed to read frame at index {i}.")
                    continue

                frame_path = f"{uid}-vid_tmp{i}.png"
                cv2.imwrite(frame_path, frame_f)

                if not os.path.exists(frame_path):
                    print(f"Error: Failed to save frame at index {i}.")
                    continue

                out_url = f'https://nymbo-reverse-image.hf.space/file={os.path.abspath(frame_path)}'
                res = rev_img_searcher.search(out_url)

                if len(res) > 0:
                    out_cnt = 0
                    for search_item in res:
                        out_cnt += 1
                        html_out += f"""
                        <div>
                        Title: {search_item.page_title}<br>
                        Site: <a href='{search_item.page_url}' target='_blank'>{search_item.page_url}</a><br>
                        Img: <a href='{search_item.image_url}' target='_blank'>{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
    except Exception as e:
        print(f"Exception during video processing: {e}")
        return gr.HTML(f"<h1>Error: {e}</h1>"), "", ""
    return gr.HTML('No frame matches found.'), "", ""

# Function to process image file
def process_im(file, url):
    if not url.startswith("https://nymbo"):
        return url
    try:
        read_file = Image.open(file)
        read_file.save(f"{uid}-tmp.png")
        out_url = f'https://nymbo-reverse-image.hf.space/file={os.path.abspath(f"{uid}-tmp.png")}'
        return out_url
    except Exception as e:
        print(f"Exception during image processing: {e}")
        return gr.HTML(f"<h1>Error: {e}</h1>")

# Function to perform reverse image search
def rev_im(image):
    if not os.path.exists(image):
        print("Error: Image file does not exist.")
        return gr.HTML("<h1>Error: Image file not found.</h1>")

    try:
        image = cv2.imread(image)
        if image is None:
            print("Error: Failed to read image.")
            return gr.HTML("<h1>Error: Could not read image.</h1>")

        cv2.imwrite(f"{uid}-im_tmp.png", image)
        out_url = f'https://nymbo-reverse-image.hf.space/file={os.path.abspath(f"{uid}-im_tmp.png")}'
        rev_img_searcher = ReverseImageSearcher()
        res = rev_img_searcher.search(out_url)

        html_out = ""
        count = 0
        for search_item in res:
            count += 1
            html_out += f"""
            <div>
            Title: {search_item.page_title}<br>
            Site: <a href='{search_item.page_url}' target='_blank'>{search_item.page_url}</a><br>
            Img: <a href='{search_item.image_url}' target='_blank'>{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}')
    except Exception as e:
        print(f"Exception during reverse image search: {e}")
        return gr.HTML(f"<h1>Error: {e}</h1>")

# Gradio app layout
with gr.Blocks() as app:
    source_tog = gr.Radio(choices=["Image", "Video"], value="Image")
    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("Search Image")
    vid_url = gr.Textbox(label="Video URL")
    vid_url_btn = gr.Button("Load Video")
    inp_vid = gr.Video(label="Search Video")
    every_n = gr.Number(label="Every nth Frame", value=10)
    go_btn_vid = gr.Button("Start Video Search")
    html_out = gr.HTML("")

    source_tog.change(lambda x: (gr.update(visible=x == "Image"), gr.update(visible=x == "Video")), [source_tog], [inp_im, inp_vid])
    load_im_btn.click(lambda url: url, inp_url, inp_im)
    go_btn_im.click(rev_im, inp_im, html_out)
    vid_url_btn.click(dl, vid_url, [inp_vid, html_out])
    go_btn_vid.click(process_vid, [inp_vid, "", every_n], [html_out])

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