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 | |
| from pathlib import Path | |
| import shutil | |
| import logging | |
| import sys | |
| # Set up logging | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format='%(asctime)s - %(levelname)s - %(message)s', | |
| handlers=[ | |
| logging.StreamHandler(sys.stdout) | |
| ] | |
| ) | |
| logger = logging.getLogger(__name__) | |
| # Custom CSS for better presentation | |
| CUSTOM_CSS = """ | |
| .result-item { | |
| margin: 20px 0; | |
| padding: 15px; | |
| border: 1px solid #ddd; | |
| border-radius: 5px; | |
| } | |
| .result-item img { | |
| max-width: 100%; | |
| height: auto; | |
| margin: 10px 0; | |
| } | |
| .error-message { | |
| color: red; | |
| padding: 10px; | |
| border: 1px solid red; | |
| border-radius: 5px; | |
| margin: 10px 0; | |
| } | |
| """ | |
| def create_temp_dir(): | |
| """Create a temporary directory with UUID.""" | |
| uid = str(uuid.uuid4()) | |
| temp_dir = Path(f"temp_{uid}") | |
| temp_dir.mkdir(exist_ok=True) | |
| return temp_dir | |
| def cleanup_temp_dir(temp_dir): | |
| """Safely clean up temporary directory and its contents.""" | |
| try: | |
| if temp_dir.exists(): | |
| shutil.rmtree(str(temp_dir)) | |
| except Exception as e: | |
| logger.error(f"Error cleaning up temporary directory: {e}") | |
| def dl(inp): | |
| """Download video from URL using yt-dlp.""" | |
| if not inp or not inp.strip(): | |
| return None, gr.HTML("<p class='error-message'>Please provide a valid URL</p>"), "", "" | |
| temp_dir = create_temp_dir() | |
| try: | |
| # Sanitize input filename | |
| inp_out = inp.replace("https://", "").replace("/", "_").replace(".", "_").replace("=", "_").replace("?", "_") | |
| output_path = str(temp_dir / f"{inp_out}.mp4") | |
| # Configure yt-dlp options | |
| ydl_opts = { | |
| 'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/mp4', | |
| 'outtmpl': output_path, | |
| 'quiet': True, | |
| 'no_warnings': True, | |
| 'extract_flat': False, | |
| 'merge_output_format': 'mp4' | |
| } | |
| # Handle Twitter URLs differently | |
| if "twitter" in inp.lower(): | |
| ydl_opts['extractor_args'] = {'twitter': {'api': ['syndication']}} | |
| # Download the video | |
| with yt_dlp.YoutubeDL(ydl_opts) as ydl: | |
| logger.info(f"Downloading video from: {inp}") | |
| ydl.download([inp]) | |
| if os.path.exists(output_path): | |
| logger.info(f"Successfully downloaded to: {output_path}") | |
| return output_path, gr.HTML(""), "", "" | |
| else: | |
| raise Exception("Download completed but file not found") | |
| except Exception as e: | |
| logger.error(f"Error downloading video: {e}") | |
| cleanup_temp_dir(temp_dir) | |
| return None, gr.HTML(f"<p class='error-message'>Error: {str(e)}</p>"), "", "" | |
| def process_vid(file, cur_frame, every_n): | |
| """Process video file and search for frames.""" | |
| if not file: | |
| return gr.HTML("<p class='error-message'>No video file provided</p>"), "", "" | |
| temp_dir = create_temp_dir() | |
| capture = None | |
| try: | |
| capture = cv2.VideoCapture(str(file)) | |
| if not capture.isOpened(): | |
| raise Exception("Failed to open video file") | |
| 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 and cur_frame.strip() else 0 | |
| for i in range(start_frame, frame_count-1): | |
| if count == int(every_n): | |
| count = 1 | |
| logger.info(f"Processing frame {i}") | |
| capture.set(cv2.CAP_PROP_POS_FRAMES, i) | |
| ret, frame = capture.read() | |
| if not ret: | |
| continue | |
| temp_image = temp_dir / f"frame_{i}.png" | |
| cv2.imwrite(str(temp_image), frame) | |
| # Generate the URL for the temporary image | |
| image_path = os.path.abspath(str(temp_image)) | |
| image_url = f'https://nymbo-reverse-image.hf.space/file={image_path}' | |
| try: | |
| results = rev_img_searcher.search(image_url) | |
| if results: | |
| out_cnt = len(results) | |
| html_out = build_results_html(results) | |
| return gr.HTML(f'<h1>Total Found: {out_cnt}</h1><br>{html_out}'), f"Found frame: {i}", i+int(every_n) | |
| except Exception as search_error: | |
| logger.error(f"Search error on frame {i}: {search_error}") | |
| count += 1 | |
| except Exception as e: | |
| logger.error(f"Error processing video: {e}") | |
| return gr.HTML(f'<p class="error-message">Error: {str(e)}</p>'), "", "" | |
| finally: | |
| if capture is not None: | |
| capture.release() | |
| cleanup_temp_dir(temp_dir) | |
| return gr.HTML('No frame matches found.'), "", "" | |
| def process_im(file, url): | |
| """Process image file or URL.""" | |
| if not url.startswith("https://nymbo"): | |
| return url | |
| temp_dir = create_temp_dir() | |
| try: | |
| temp_image = temp_dir / "temp.png" | |
| read_file = Image.open(file) | |
| read_file.save(str(temp_image)) | |
| out_url = f'https://nymbo-reverse-image.hf.space/file={os.path.abspath(str(temp_image))}' | |
| return out_url | |
| finally: | |
| cleanup_temp_dir(temp_dir) | |
| def rev_im(image): | |
| """Reverse image search for a single image.""" | |
| if not image: | |
| return gr.HTML("<p class='error-message'>No image provided</p>") | |
| temp_dir = create_temp_dir() | |
| try: | |
| temp_image = temp_dir / "temp.png" | |
| image_cv = cv2.imread(image) | |
| cv2.imwrite(str(temp_image), image_cv) | |
| out_url = f'https://nymbo-reverse-image.hf.space/file={os.path.abspath(str(temp_image))}' | |
| rev_img_searcher = ReverseImageSearcher() | |
| results = rev_img_searcher.search(out_url) | |
| return gr.HTML(build_results_html(results, len(results))) | |
| except Exception as e: | |
| logger.error(f"Error in reverse image search: {e}") | |
| return gr.HTML(f'<p class="error-message">Error: {str(e)}</p>') | |
| finally: | |
| cleanup_temp_dir(temp_dir) | |
| def build_results_html(results, count=None): | |
| """Build HTML output for search results.""" | |
| html_out = f"<h1>Total Found: {count}</h1><br>" if count is not None else "" | |
| for item in results: | |
| html_out += f""" | |
| <div class="result-item"> | |
| <h3>{item.page_title}</h3> | |
| <p>Source: <a href='{item.page_url}' target='_blank' rel='noopener noreferrer'>{item.page_url}</a></p> | |
| <p>Image: <a href='{item.image_url}' target='_blank' rel='noopener noreferrer'>{item.image_url}</a></p> | |
| <img class='my_im' src='{item.image_url}' alt='{item.page_title}'> | |
| </div> | |
| """ | |
| return html_out | |
| def create_ui(): | |
| """Create the Gradio interface.""" | |
| with gr.Blocks(css=CUSTOM_CSS) as app: | |
| with gr.Row(): | |
| gr.Column() | |
| with gr.Column(): | |
| source_tog = gr.Radio( | |
| choices=["Image", "Video"], | |
| value="Image", | |
| label="Search Type" | |
| ) | |
| # Image search interface | |
| with gr.Box(visible=True) as im_box: | |
| inp_url = gr.Textbox(label="Image URL") | |
| load_im_btn = gr.Button("Load Image", variant="primary") | |
| inp_im = gr.Image(label="Search Image", type='filepath') | |
| go_btn_im = gr.Button("Search", variant="primary") | |
| # Video search interface | |
| with gr.Box(visible=False) as vid_box: | |
| vid_url = gr.Textbox(label="Video URL") | |
| vid_url_btn = gr.Button("Load URL", variant="primary") | |
| inp_vid = gr.Video(label="Search Video") | |
| with gr.Row(): | |
| every_n = gr.Number( | |
| label="Every nth frame", | |
| value=10, | |
| minimum=1, | |
| maximum=100, | |
| step=1 | |
| ) | |
| stat_box = gr.Textbox(label="Status", interactive=False) | |
| with gr.Row(): | |
| go_btn_vid = gr.Button("Start", variant="primary") | |
| next_btn = gr.Button("Next Frame", variant="secondary") | |
| gr.Column() | |
| with gr.Row(): | |
| html_out = gr.HTML() | |
| with gr.Row(visible=False): | |
| hid_box = gr.Textbox() | |
| # Event handlers | |
| def toggle_interface(tog): | |
| """Toggle between image and video interfaces.""" | |
| return ( | |
| gr.update(visible=tog == "Image"), | |
| gr.update(visible=tog == "Video") | |
| ) | |
| # Connect all event handlers | |
| source_tog.change( | |
| toggle_interface, | |
| [source_tog], | |
| [im_box, vid_box], | |
| cancels=[go_btn_vid.click, go_btn_im.click, load_im_btn.click, vid_url_btn.click] | |
| ) | |
| load_im_btn.click(lambda x: x, inp_url, inp_im) | |
| next_btn.click(process_vid, [inp_vid, hid_box, every_n], [html_out, stat_box, hid_box]) | |
| vid_url_btn.click(dl, vid_url, [inp_vid, html_out, stat_box, hid_box]) | |
| go_btn_vid.click(process_vid, [inp_vid, hid_box, every_n], [html_out, stat_box, hid_box]) | |
| go_btn_im.click(rev_im, inp_im, html_out) | |
| return app | |
| if __name__ == "__main__": | |
| app = create_ui() | |
| app.queue(concurrency_count=20).launch() |