Update app.py

app.py

@@ -1,173 +1,226 @@
 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
+from pathlib import Path
 
-# Generate a unique identifier for temporary file storage
-uid = str(uuid.uuid4())
+def ensure_directory(directory):
+    Path(directory).mkdir(parents=True, exist_ok=True)
+    return directory
 
-# Ensure the temporary directory exists
-if not os.path.exists(uid):
-    os.makedirs(uid)
+# Create a unique working directory for each session
+def get_session_dir():
+    session_id = str(uuid.uuid4())
+    return ensure_directory(f"temp_{session_id}")
 
-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):
+    if not inp or not inp.strip():
+        return None, gr.HTML("Please provide a valid URL"), "", ""
+    
+    work_dir = get_session_dir()
     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')
+        # Sanitize input filename
+        inp_out = inp.replace("https://", "").replace("/", "_").replace(".", "_").replace("=", "_").replace("?", "_")
+        output_path = os.path.join(work_dir, f"{inp_out}.mp4")
+        
+        ydl_opts = {
+            'format': 'best[ext=mp4]',
+            'outtmpl': output_path,
+            'quiet': True,
+            'no_warnings': True
+        }
+        
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            ydl.download([inp])
+        
+        if os.path.exists(output_path):
+            out = output_path
+            print(f"Successfully downloaded video to: {out}")
         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}")
+            raise Exception("Video download failed")
+            
     except Exception as e:
-        print(f"Exception during video download: {e}")
-        return None, gr.HTML(f"<h1>Error: {e}</h1>"), "", ""
+        print(f"Download error: {str(e)}")
+        return None, gr.HTML(f"Error downloading video: {str(e)}"), "", ""
+        
     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 file or not os.path.exists(file):
+        return gr.HTML("No valid video file provided."), "", ""
+    
+    work_dir = os.path.dirname(file)
+    capture = cv2.VideoCapture(file)
+    
     if not capture.isOpened():
-        print("Error: Video file could not be opened.")
-        return gr.HTML("<h1>Error: Failed to open video file.</h1>"), "", ""
-
+        return gr.HTML("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 else 0
-
+    count = int(every_n) if every_n else 10
+    start_frame = int(cur_frame) if cur_frame and cur_frame.strip() else 0
+    
     try:
-        for i in range(start_frame, frame_count - 1):
+        for i in range(start_frame, frame_count-1):
             if count == int(every_n):
                 count = 1
+                print(f"Processing frame {i}")
+                
+                # Read frame
                 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}.")
+                ret, frame = capture.read()
+                
+                if not ret or frame is None:
+                    print(f"Failed to read frame {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}.")
+                
+                # Save frame
+                frame_path = os.path.join(work_dir, f"frame_{i}.png")
+                if not cv2.imwrite(frame_path, frame):
+                    print(f"Failed to write frame {i}")
                     continue
-
+                
+                # Process frame
                 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:
+                
+                if res:
                     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>
+                            <h3>Result {out_cnt}</h3>
+                            <p>Title: {search_item.page_title}</p>
+                            <p>Site: <a href='{search_item.page_url}' target='_blank'>{search_item.page_url}</a></p>
+                            <p>Image: <a href='{search_item.image_url}' target='_blank'>{search_item.image_url}</a></p>
+                            <img class='my_im' src='{search_item.image_url}' alt='Search result'>
                         </div>
                         """
-                    return gr.HTML(f'<h1>Total Found: {out_cnt}</h1><br>{html_out}'), f"Found frame: {i}", i + int(every_n)
+                    return gr.HTML(f'<h2>Total Found: {out_cnt}</h2>{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.'), "", ""
+        import traceback
+        error_msg = f"Error processing video: {str(e)}\n{traceback.format_exc()}"
+        print(error_msg)
+        return gr.HTML(error_msg), "", ""
+    finally:
+        capture.release()
+        
+    return gr.HTML('No matches found in processed frames.'), "", ""
 
-# Function to process image file
 def process_im(file, url):
     if not url.startswith("https://nymbo"):
         return url
+    
+    work_dir = get_session_dir()
+    output_path = os.path.join(work_dir, "processed_image.png")
+    
     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
+        read_file.save(output_path)
+        return f'https://nymbo-reverse-image.hf.space/file={os.path.abspath(output_path)}'
     except Exception as e:
-        print(f"Exception during image processing: {e}")
-        return gr.HTML(f"<h1>Error: {e}</h1>")
+        print(f"Error processing image: {str(e)}")
+        return url
 
-# 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>")
-
+    if not image or not os.path.exists(image):
+        return gr.HTML("No valid image provided.")
+    
+    work_dir = get_session_dir()
+    html_out = ""
+    
     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")}'
+        # Read and write image
+        img = cv2.imread(image)
+        if img is None:
+            return gr.HTML("Failed to read image file.")
+            
+        output_path = os.path.join(work_dir, "search_image.png")
+        if not cv2.imwrite(output_path, img):
+            return gr.HTML("Failed to process image file.")
+            
+        # Search image
+        out_url = f'https://nymbo-reverse-image.hf.space/file={os.path.abspath(output_path)}'
         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>
+                <h3>Result {count}</h3>
+                <p>Title: {search_item.page_title}</p>
+                <p>Site: <a href='{search_item.page_url}' target='_blank'>{search_item.page_url}</a></p>
+                <p>Image: <a href='{search_item.image_url}' target='_blank'>{search_item.image_url}</a></p>
+                <img class='my_im' src='{search_item.image_url}' alt='Search result'>
             </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>")
+        return gr.HTML(f"Error processing image: {str(e)}")
+        
+    return gr.HTML(f'<h2>Total Found: {count}</h2>{html_out}')
 
-# Gradio app layout
+# Gradio Interface
 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)
+    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")
+                inp_im = gr.Image(label="Search Image", type='filepath')
+                go_btn_im = gr.Button("Search")
+            
+            # 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")
+                inp_vid = gr.Video(label="Search Video")
+                with gr.Row():
+                    every_n = gr.Number(label="Process every N frames", value=10, minimum=1)
+                stat_box = gr.Textbox(label="Status")
+                with gr.Row():
+                    go_btn_vid = gr.Button("Start Search")
+                    next_btn = gr.Button("Next Frame")
+                
+        gr.Column()
+        
+    with gr.Row():
+        html_out = gr.HTML()
+    with gr.Row(visible=False):
+        hid_box = gr.Textbox()
+
+    # Interface logic
+    def toggle_interface(tog):
+        return gr.update(visible=tog == "Image"), gr.update(visible=tog == "Video")
+        
+    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]
+    )
+    
+    # Button actions
+    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)
-    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()