Update app.py

app.py

@@ -1,48 +1,35 @@
 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 = str(uuid.uuid4())
+# Unique identifier for files
+uid = uuid.uuid4()
 
+# JavaScript for displaying image dimensions
 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 imgSize(){
+        var myImg = document.getElementsByClassName("my_im");
+        var realWidth = myImg.naturalWidth;
+        var realHeight = myImg.naturalHeight;
+        alert("Original width=" + realWidth + ", " + "Original height=" + realHeight);
+    }
 """
 
 def dl(inp):
     out = None
     try:
-        inp_out = inp.replace("https://", "")
-        inp_out = inp_out.replace("/", "_").replace(".", "_").replace("=", "_").replace("?", "_")
-        output_path = f'{uid}/{inp_out}.mp4'
-
-        ydl_opts = {
-            'outtmpl': output_path,
-            'format': 'mp4',
-            'recodevideo': 'mp4',
-            'trim_file_name': 160,
-            'postprocessors': [{
-                'key': 'FFmpegVideoConvertor',
-                'preferedformat': 'mp4',
-            }],
-        }
-
+        inp_out = inp.replace("https://", "").replace("/", "_").replace(".", "_").replace("=", "_").replace("?", "_")
         if "twitter" in inp:
-            ydl_opts['extractor_args'] = {'twitter': {'api': 'syndication'}}
-
-        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
-            ydl.download([inp])
-
-        out = output_path
+            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"
         print(out)
     except Exception as e:
         print(e)
@@ -53,36 +40,33 @@ def process_vid(file, cur_frame, every_n):
     if not os.path.isfile(new_video_in):
         print(f"Video file {new_video_in} does not exist.")
         return (gr.HTML('Video file not found.'), "", "")
-
+    
     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)
-
+    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
-                print(f"Processing frame {i}")
+                print(i)
                 capture.set(cv2.CAP_PROP_POS_FRAMES, i)
                 ret, frame_f = capture.read()
                 if not ret:
                     print(f"Failed to read frame {i}")
                     continue
-                frame_filename = f"{uid}-vid_tmp{i}.png"
-                cv2.imwrite(frame_filename, frame_f)
-                out = os.path.abspath(frame_filename)
+                
+                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(f"Frame saved at {out}")
-
+                print(out)
+                
                 res = rev_img_searcher.search(out_url)
                 out_cnt = 0
-                if len(res) > 0:
+                if res:
                     for search_item in res:
                         out_cnt += 1
                         html_out += f"""
@@ -92,12 +76,13 @@ def process_vid(file, cur_frame, every_n):
                         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))
-            else:
-                count += 1
+                    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)
     except Exception as e:
-        return (gr.HTML(f'{e}'), "", "")
-    return (gr.HTML('No frame matches found.'), "", "")
+        return gr.HTML(f'{e}'), "", ""
+    
+    return gr.HTML('No frame matches found.'), "", ""
 
 def process_im(file, url):
     if not url.startswith("https://nymbo"):
@@ -113,6 +98,9 @@ def process_im(file, url):
 def rev_im(image):
     html_out = ""
     image = cv2.imread(image)
+    if image is None:
+        return gr.HTML("Failed to load 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}'    
@@ -128,6 +116,7 @@ def rev_im(image):
         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:
@@ -139,7 +128,7 @@ with gr.Blocks() as app:
                 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")
+                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")
@@ -150,6 +139,7 @@ with gr.Blocks() as app:
                 with gr.Row():
                     go_btn_vid = gr.Button("Start")
                     next_btn = gr.Button("Next")
+                
         gr.Column()
     with gr.Row():
         html_out = gr.HTML("")
@@ -157,14 +147,11 @@ with gr.Blocks() as app:
         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)
-
+        return (gr.update(visible=True), gr.update(visible=False)) if tog == "Image" else (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])