Update app.py

app.py

@@ -1,178 +1,285 @@
 import gradio as gr
-import requests
+import requests 
 import yt_dlp
 import cv2
 from google_img_source_search import ReverseImageSearcher
 from PIL import Image
-import os
+import os 
 import uuid
-import subprocess
+from pathlib import Path
+import shutil
+import logging
+import sys
 
-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);
-    }"""
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler(sys.stdout)
+    ]
+)
+logger = logging.getLogger(__name__)
 
-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
+# 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;
+}
+"""
 
-        # Generate a unique output filename
-        inp_out = inp.replace("https://", "").replace("/", "_").replace(".", "_").replace("=", "_").replace("?", "_")
+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
 
-        # Download using yt_dlp
-        command = [
-            "yt-dlp",
-            inp,
-            "--trim-filenames", "160",
-            "-o", f"{uid}/{inp_out}.mp4",
-            "-S", "res,mp4",
-            "--recode", "mp4"
-        ]
+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}")
 
-        if "twitter" in inp:
-            command.insert(2, "--extractor-arg")
-            command.insert(3, "twitter:api=syndication")
+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>"), "", ""
 
-        # Use subprocess to execute the command and handle errors
-        result = subprocess.run(command, capture_output=True, text=True)
+    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")
 
-        # 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>"), "", ""
+        # 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'
+        }
 
-        out = f"{uid}/{inp_out}.mp4"
-        print(out)
+        # 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:
-        print(e)
-        return None, gr.HTML(f"<p>Error: {str(e)}</p>"), "", ""
-
-    return out, gr.HTML(""), "", ""
+        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):
-    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)
+    """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:
-        for i in range(start_frame, frame_count - 1):
+        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
-                print(i)
+                logger.info(f"Processing frame {i}")
                 capture.set(cv2.CAP_PROP_POS_FRAMES, i)
-                ret, frame_f = capture.read()
+                ret, frame = 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)
+                    
+                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
-            print(i + 1)
+            
     except Exception as e:
-        return gr.HTML(f"{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
-    else:
+    
+    temp_dir = create_temp_dir()
+    try:
+        temp_image = temp_dir / "temp.png"
         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}'
+        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):
-    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)
+    """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 ""
     
-    def load_image(url):
-        return url
+    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)
 
-    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])
+    return app
 
-app.queue(concurrency_count=20).launch()
+if __name__ == "__main__":
+    app = create_ui()
+    app.queue(concurrency_count=20).launch()