Update app.py

app.py

@@ -6,7 +6,9 @@ from google_img_source_search import ReverseImageSearcher
 from PIL import Image
 import os 
 import uuid
-uid = str(uuid.uuid4())
+
+uid = uuid.uuid4()
+
 size_js = """
     function imgSize(){
         var myImg = document.getElementsByClassName("my_im");
@@ -19,8 +21,12 @@ 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("?", "_")
+        
         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')  
         else:
@@ -33,32 +39,56 @@ def dl(inp):
     return out, gr.HTML(""), "", ""
 
 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)
     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)
+
+        # 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):
             if count == int(every_n):
                 count = 1
+                print(i)
                 capture.set(cv2.CAP_PROP_POS_FRAMES, i)
                 ret, frame_f = capture.read()
-                if not ret:
-                    print(f"Failed to capture frame at {i}")
+                
+                if not ret or frame_f is None:
                     continue
-                cv2.imwrite(f"{uid}-vid_tmp{i}.png", frame_f)
-                out = os.path.abspath(f"{uid}-vid_tmp{i}.png")
+                    
+                temp_img_path = f"{uid}/vid_tmp{i}.png"
+                cv2.imwrite(temp_img_path, frame_f)
+                
+                if not os.path.exists(temp_img_path):
+                    continue
+                    
+                out = os.path.abspath(temp_img_path)
                 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:
+                        print(f'counting {count}')
                         out_cnt += 1
                         html_out = f"""{html_out}
                         <div>
@@ -67,56 +97,77 @@ 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>"""
+                    
+                    capture.release()
                     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()
+        
     except Exception as e:
-        return (gr.HTML(f'{e}'), "", "")
+        return (gr.HTML(f'Error: {str(e)}'), "", "")
+        
     return (gr.HTML('No frame matches found.'), "", "")
 
 def process_im(file, url):
     if not url.startswith("https://nymbo"):
         return url
     else:
-        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)
+        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
 
 def rev_im(image):
-    out_list = []
-    out_im = []
-    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
-        out_dict = {
-            'Title': f'{search_item.page_title}',
-            'Site': f'{search_item.page_url}',
-            'Img': f'{search_item.image_url}',
-        }
-        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>"""
+    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)
         
-    return (gr.HTML(f'<h1>Total Found: {count}</h1><br>{html_out}'))
+        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}'    
+        
+        rev_img_searcher = ReverseImageSearcher()
+        res = rev_img_searcher.search(out_url)
+        count = 0
+        html_out = ""
+        
+        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}')
+        
+    except Exception as e:
+        return gr.HTML(f"Error: {str(e)}")
 
 with gr.Blocks() as app:
     with gr.Row():
         gr.Column()
         with gr.Column():
-            source_tog = gr.Radio(choices=["Image","Video"], value="Image")
+            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")
@@ -132,18 +183,22 @@ with gr.Blocks() as app:
                 with gr.Row():
                     go_btn_vid = gr.Button("Start")
                     next_btn = gr.Button("Next")
-        gr.Column()            
+        gr.Column()
+    
     with gr.Row():
-        html_out = gr.HTML("""""")
+        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)    
+            
     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])