Update app.py

app.py

@@ -4,6 +4,8 @@ import os
 import json
 import rembg
 import cv2
+import uuid
+uid=uuid.uuid4()
 
 load_js = """
 function(text_input, url_params) {
@@ -26,8 +28,43 @@ def rem_cv(inp):
         fgmask = cv2.morphologyEx(fgmask, cv2.MORPH_OPEN, kernel)
         yield (fgmask)
 
+def load_video(vid):
+    new_video_in = str(vid)
+    capture = cv2.VideoCapture(new_video_in)
+    frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
+    tmp_frame='1'
+    capture.release()
+    return frame_count 
+
+
+
+def capture_function(vid=None):
+    if vid==None:
+        vid=img_file
+    new_video_in = str(vid)
+    capture = cv2.VideoCapture(new_video_in)    
+    fbox=[]
+    cnt=0
+    frame_count1= int(frame_count)
+    for i in range(int(frame_count1)):
+        capture.set(cv2.CAP_PROP_POS_FRAMES, i)
+        ret, frame_f = capture.read(i)
+        frame_ff = cv2.cvtColor(frame_f, cv2.COLOR_BGR2RGB)
         
+        fbox.append(frame_ff)
+
+    capture.set(cv2.CAP_PROP_POS_FRAMES, int(frame_count)-1)
+    ret, frame_f = capture.read()
+    frame_ff = cv2.cvtColor(frame_f, cv2.COLOR_BGR2RGB)
+    fbox.append(frame_ff)
+    capture.release()
+    return fbox
+
+    
 def rem_bg(inp):
+    cap = cv2.VideoCapture(f'{inp}')
+    ret, frame = cap.read()
+    
     out = rem_bg()
 
 
@@ -76,7 +113,12 @@ with gr.Blocks() as app:
             with gr.Column():
                 outp_vid=gr.Video()
             with gr.Column():
-                outp_file=gr.Textbox()
+                frame_count=gr.Textbox()
+                outp_file=gr.Files()
+    with gr.Tab("Frames"):
+        frame_btn = gr.Button("Get Frames")                
+        frame_gal = gr.Gallery()
+        
     with gr.Tab("Rem BG"):
         with gr.Row():    
             with gr.Column():
@@ -84,11 +126,13 @@ with gr.Blocks() as app:
                 in_vid=gr.Video()
             with gr.Column():
                 #rem_vid=gr.Video()
-                rem_vid=gr.Image()
+                rem_vid=gr.Gallery()
     with gr.Row(visible=False):
         text_input=gr.Textbox()
         text_output=gr.Textbox()
         url_params=gr.JSON()
+    outp_vid.change(load_video,outp_vid,frame_count)
+    frame_btn.click(capture_function,[outp_vid,frame_count],frame_gal)
     rem_btn.click(rem_cv,in_vid,rem_vid)
     go_btn.click(dl,inp_url,[outp_vid,in_vid,outp_file])
     app.load(fn=predict, inputs=[text_input,url_params], outputs=[text_output,text_input],_js=load_js)