Update app.py

app.py

@@ -6,8 +6,11 @@ import rembg
 import cv2
 import uuid
 uid=uuid.uuid4()
-if not os.path.exists(f'{uid}'): os.makedirs(f'{uid}')
+if not os.path.exists(f'{uid}-frames'): os.makedirs(f'{uid}')
+if not os.path.exists(f'{uid}-rembg'): os.makedirs(f'{uid}')
+
 load_js = """
+
 function(text_input, url_params) {
     console.log(text_input, url_params);
     const params = new URLSearchParams(window.location.search);
@@ -39,10 +42,10 @@ def load_video(vid):
 
 
 def capture_function(vid=None,frame_count=None):
-    if vid==None:
-        vid=img_file
     new_video_in = str(vid)
-    capture = cv2.VideoCapture(new_video_in)    
+    capture = cv2.VideoCapture(new_video_in)
+    #frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
+    
     fbox=[]
     cnt=0
     frame_count1= int(frame_count)
@@ -53,21 +56,40 @@ def capture_function(vid=None,frame_count=None):
         cv2.imwrite(f'{uid}/{i}.png',frame_ff)
         fbox.append(f'{uid}/{i}.png')
         
-    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)
-    cv2.imwrite(f'{uid}/{int(frame_count)-1}.png',frame_ff)
-    fbox.append(f'{uid}/{int(frame_count)-1}.png')
-    capture.release()
+    #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)
+    #cv2.imwrite(f'{uid}/{int(frame_count)-1}.png',frame_ff)
+    #fbox.append(f'{uid}/{int(frame_count)-1}.png')
+    #capture.release()
     return fbox,fbox
 
     
-def rem_bg(inp):
-    cap = cv2.VideoCapture(f'{inp}')
-    ret, frame = cap.read()
+def rem_bg(vid,frame_count):
+    new_video_in = str(vid)
+    capture = cv2.VideoCapture(new_video_in)    
+    fbox2=[]
+    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)
+        out = rem_bg(frame_ff)
+
+        cv2.imwrite(f'{uid}-rembg/{i}.png',out)
+        fbox2.append(f'{uid}-rembg/{i}.png')
+        
+    #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)
+    #cv2.imwrite(f'{uid}/{int(frame_count)-1}.png',frame_ff)
+    #fbox.append(f'{uid}/{int(frame_count)-1}.png')
+    #capture.release()
+    #out = cv2.VideoWriter('video.avi',cv2.VideoWriter_fourcc(*'DIVX'), 15, size)
+    return fbox2
+
     
-    out = rem_bg()
-    out = cv2.VideoWriter('video.avi',cv2.VideoWriter_fourcc(*'DIVX'), 15, size)
 
 
     
@@ -121,7 +143,7 @@ with gr.Blocks() as app:
         frame_btn = gr.Button("Get Frames")
         with gr.Row():
             with gr.Column():
-                frame_gal = gr.Gallery(columns=6,show_label=True)
+                frame_gal = gr.Gallery(columns=6)
             with gr.Column():
                 frame_file = gr.Files()
         
@@ -132,14 +154,14 @@ with gr.Blocks() as app:
                 in_vid=gr.Video()
             with gr.Column():
                 #rem_vid=gr.Video()
-                rem_vid=gr.Gallery()
+                rem_vid=gr.Gallery(columns=6)
     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,frame_file])
-    rem_btn.click(rem_cv,in_vid,rem_vid)
+    rem_btn.click(rem_cv,[outp_vid,frame_count],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)
 app.launch()