Update app.py

app.py

@@ -1,14 +1,13 @@
 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
 
 uid = uuid.uuid4()
-
 size_js = """
     function imgSize(){
         var myImg = document.getElementsByClassName("my_im");
@@ -17,193 +16,158 @@ size_js = """
         alert("Original width=" + realWidth + ", " + "Original height=" + realHeight);
     }"""
 
+# Function to download video from URL
 def dl(inp):
     out = None
-    out_file = []
     try:
-        # Create the directory if it doesn't exist
-        os.makedirs(str(uid), exist_ok=True)
-        
-        inp_out = inp.replace("https://", "")
-        inp_out = inp_out.replace("/", "_").replace(".", "_").replace("=", "_").replace("?", "_")
-        
+        inp_out = inp.replace("https://", "").replace("/", "_").replace(".", "_").replace("=", "_").replace("?", "_")
+        video_path = f"{uid}/{inp_out}.mp4"
+
         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')  
+            os.system(f'yt-dlp "{inp}" --extractor-arg "twitter:api=syndication" --trim-filenames 160 -o "{video_path}" -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"
-        print(out)
+            os.system(f'yt-dlp "{inp}" --trim-filenames 160 -o "{video_path}" -S res,mp4 --recode mp4')
+
+        # Verify video download
+        if not os.path.exists(video_path):
+            print(f"Downloaded video {video_path} does not exist.")
+            return None, gr.HTML("Download failed."), "", ""
+        if os.path.getsize(video_path) == 0:
+            print(f"Downloaded video {video_path} is empty.")
+            return None, gr.HTML("Downloaded video is empty."), "", ""
+
+        out = video_path
+        print(f"Video downloaded successfully: {out}")
     except Exception as e:
-        print(e)
+        print(f"Error downloading video: {e}")
     return out, gr.HTML(""), "", ""
 
+# Function to process video frames
 def process_vid(file, cur_frame, every_n):
-    try:
-        # Ensure the file exists and is readable
-        if not os.path.exists(file):
-            return gr.HTML("Error: Video file not found."), "", ""
-            
-        new_video_in = str(file)
-        capture = cv2.VideoCapture(new_video_in)
-        
-        if not capture.isOpened():
-            return gr.HTML("Error: Could not open video file."), "", ""
-            
-        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)
+    if not file:
+        return gr.HTML("No video file provided."), "", ""
+
+    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)
+    start_frame = int(cur_frame) if cur_frame else 0
 
-        # Create directory for temporary files if it doesn't exist
-        os.makedirs(str(uid), exist_ok=True)
-            
-        for i in range(start_frame, frame_count-1):
+    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()
-                
+
+                # Validate frame read
                 if not ret or frame_f is None:
+                    print(f"Failed to read frame at index {i}.")
                     continue
-                    
-                temp_img_path = f"{uid}/vid_tmp{i}.png"
-                cv2.imwrite(temp_img_path, frame_f)
-                
-                if not os.path.exists(temp_img_path):
+
+                # Check if frame is empty
+                if frame_f.size == 0:
+                    print(f"Frame at index {i} is empty.")
                     continue
-                    
-                out = os.path.abspath(temp_img_path)
+
+                frame_path = f"{uid}-vid_tmp{i}.png"
+                cv2.imwrite(frame_path, frame_f)
+                out = os.path.abspath(frame_path)
                 out_url = f'https://nymbo-reverse-image.hf.space/file={out}'
-                print(out)
-                
+                print(f"Frame saved: {out}")
+
+                # Perform reverse image search
                 res = rev_img_searcher.search(out_url)
-                out_cnt = 0
-                
-                if len(res) > 0:
+                if res:
+                    out_cnt = 0
                     for search_item in res:
-                        print(f'counting {count}')
                         out_cnt += 1
-                        html_out = f"""{html_out}
+                        html_out += f"""
                         <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>
+                        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>"""
-                    
-                    capture.release()
-                    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'<h1>Total Found: {out_cnt}</h1><br>{html_out}'), f"Found frame: {i}", i + int(every_n)
             count += 1
-            print(i+1)
-            
-        capture.release()
-        
+            print(f"Processed frame: {i + 1}")
     except Exception as e:
-        return (gr.HTML(f'Error: {str(e)}'), "", "")
-        
-    return (gr.HTML('No frame matches found.'), "", "")
+        return gr.HTML(f'Error during video processing: {e}'), "", ""
+    return gr.HTML('No frame matches found.'), "", ""
 
+# Function to process image input
 def process_im(file, url):
     if not url.startswith("https://nymbo"):
         return url
-    else:
-        try:
-            read_file = Image.open(file)
-            os.makedirs(str(uid), exist_ok=True)
-            output_path = f"{uid}/tmp.png"
-            read_file.save(output_path)
-            out = os.path.abspath(output_path)
-            out_url = f'https://nymbo-reverse-image.hf.space/file={out}'
-            return out_url
-        except Exception as e:
-            print(f"Error in process_im: {str(e)}")
-            return url
+    try:
+        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
+    except Exception as e:
+        print(f"Error processing image: {e}")
+        return None
 
+# Function to perform reverse image search
 def rev_im(image):
     try:
-        os.makedirs(str(uid), exist_ok=True)
-        output_path = f"{uid}/im_tmp.png"
-        
         image = cv2.imread(image)
-        if image is None:
-            return gr.HTML("Error: Could not read image file.")
-            
-        cv2.imwrite(output_path, image)
-        
-        if not os.path.exists(output_path):
-            return gr.HTML("Error: Failed to save processed image.")
-            
-        out = os.path.abspath(output_path)
-        out_url = f'https://nymbo-reverse-image.hf.space/file={out}'    
-        
+        if image is None or image.size == 0:
+            print("Error: Image is empty or invalid.")
+            return gr.HTML("Error: Invalid 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
         html_out = ""
-        
+        count = 0
         for search_item in res:
             count += 1
-            html_out = f"""{html_out}
+            html_out += f"""
             <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>
+            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:
-        return gr.HTML(f"Error: {str(e)}")
+        print(f"Error during reverse image search: {e}")
+        return gr.HTML(f"Error: {e}")
 
+# Gradio app interface
 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()
-            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")
-                    next_btn = gr.Button("Next")
-        gr.Column()
-    
-    with gr.Row():
-        html_out = gr.HTML("")
-    with gr.Row(visible=False):
-        hid_box = gr.Textbox()
-        
+    source_tog = gr.Radio(choices=["Image", "Video"], value="Image")
+    im_box = gr.Box(visible=True)
+    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()
+    vid_box = gr.Box(visible=False)
+    vid_url = gr.Textbox(label="Video URL")
+    vid_url_btn = gr.Button("Load URL")
+    inp_vid = gr.Video(label="Search Video")
+    every_n = gr.Number(label="Every /nth frame", value=10)
+    stat_box = gr.Textbox(label="Status")
+    go_btn_vid = gr.Button("Start")
+    next_btn = gr.Button("Next")
+    html_out = gr.HTML("")
+
     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])
+        return (gr.update(visible=(tog == "Image")), gr.update(visible=(tog == "Video")))
+
+    im_load = load_im_btn.click(lambda url: url, inp_url, inp_im)
+    vid_load = vid_url_btn.click(dl, vid_url, [inp_vid, html_out, stat_box])
+    vid_proc = go_btn_vid.click(process_vid, [inp_vid, "", every_n], [html_out, stat_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])
+    source_tog.change(shuf, [source_tog], [im_box, vid_box])
 
 app.queue(concurrency_count=20).launch()