Update app.py

app.py

@@ -60,10 +60,25 @@ def capture_function(vid=None):
     return fbox,fbox
 
     
+
+def im_2_vid(images,fps):
+    #width, height = Image.open(images[0]).size
+    width, height = cv2.imread(images[0]).shape()
+    print (width,height)
+    size = (width, height)
+    codec = cv2.VideoWriter_fourcc(*'mp4v') #DIVX, DIVD
+    video = cv2.VideoWriter(vidName, codec, fps, size)
+    
+    #Do not open images in 'Image' as well as 'CV2', else following error occurs
+    #<built-in function imread> returned NULL without setting an error
+    for img in images:
+            video.write(img)
+    return (video)
 def rem_bg(vid):
     new_video_in = str(vid)
     capture = cv2.VideoCapture(new_video_in)    
     frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = capture.get(cv2.CAP_PROP_FPS)
     
     fbox2=[]
     cnt=0
@@ -77,32 +92,10 @@ def rem_bg(vid):
         cv2.imwrite(f'{uid}-rembg/{i}.png',out)
         fbox2.append(f'{uid}-rembg/{i}.png')
         frame_num=f'Working on {i+1} of {frame_count}'
-        yield fbox2,frame_num
-
-
-def im_2_vid(images,fps):
-    #width, height = Image.open(images[0]).size
-    width, height = cv2.imread(images[0]).shape()
-    print (width,height)
-    size = (width, height)
-    codec = cv2.VideoWriter_fourcc(*'mp4v') #DIVX, DIVD
-    if (len(sys.argv) < 2 and n > 0):
-        video = cv2.VideoWriter(vidName, codec, nfs, (W, H))
-    else:
-        video = cv2.VideoWriter(vidName, codec, nfs, size)
-    
-    #Do not open images in 'Image' as well as 'CV2', else following error occurs
-    #<built-in function imread> returned NULL without setting an error
-    for img in images:
-        #In case images are of different size, resize them to user specified values
-        if (len(sys.argv) < 2 and n > 0):
-            newImg = cv2.imread(img)
-            newImg = cv2.resize(newImg, (W, H))
-            video.write(newImg)
-        else:
-            video.write(img)
-   
+        yield fbox2,frame_num,None
 
+    out_vid = im_2_vid(fbox2, fps)
+    return (fbox2,frame_num,out_vid)
 
     
 def rem_bg_og(inp):
@@ -139,8 +132,28 @@ def dl(inp):
     except Exception as e:
         print (e)
         out = None
-    return out,out,out
+    return out,out,out,out
+def trim_clip(inp):
+    new_video_in = str(vid)
+    capture = cv2.VideoCapture(new_video_in)    
+    frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = capture.get(cv2.CAP_PROP_FPS)
+    fbox3=[]
+    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 = rm(frame_f)
+        print (i)
+        cv2.imwrite(f'{uid}-rembg/{i}.png',out)
+        fbox3.append(f'{uid}-rembg/{i}.png')
+        frame_num=f'Working on {i+1} of {frame_count}'
+        yield fbox2,frame_num,None    
 
+
+
+    
 with gr.Blocks() as app:
     with gr.Tab("Load"):
         
@@ -154,6 +167,14 @@ with gr.Blocks() as app:
                     frame_count=gr.Textbox(label="Frame Count",interactive = False)
                     fps=gr.Textbox(label="FPS",interactive = False)
                 outp_file=gr.Files()
+    with gr.Tab("Trim"):
+        with gr.Row():
+            start_f=gr.Number(value=0)
+            end_f=gr.Number(value=1)
+            trim_btn=gr.Button("Trim")
+        with gr.Row():
+            in_trim = gr.Video()
+            out_trim=gr.Video()
     with gr.Tab("Frames"):
         frame_btn = gr.Button("Get Frames")
         with gr.Row():
@@ -175,9 +196,11 @@ with gr.Blocks() as app:
         text_input=gr.Textbox()
         text_output=gr.Textbox()
         url_params=gr.JSON()
+    in_trim.change(None,in_trim,out_trim)
+    trim_btn.click(trim_clip,in_trim,out_trim)
     outp_vid.change(load_video,outp_vid,[frame_count,fps])
     frame_btn.click(capture_function,[outp_vid],[frame_gal,frame_file])
     rem_btn.click(rem_bg,[outp_vid],[rem_vid,frame_num])
-    go_btn.click(dl,inp_url,[outp_vid,in_vid,outp_file])
+    go_btn.click(dl,inp_url,[outp_vid,in_vid,outp_file,in_trim])
     app.load(fn=predict, inputs=[text_input,url_params], outputs=[text_output,text_input],_js=load_js)
 app.launch()