gradio state added for frame lists

app.py

@@ -27,15 +27,9 @@ sys.path.append(os.path.abspath(join(project_name, 'FGT_codes', 'FGT', 'checkpoi
 sys.path.append(os.path.abspath(join(project_name, 'FGT_codes', 'LAFC',
                 'flowCheckPoint', 'raft-things.pth')))
 
-# sys.path.append(join(project_name, 'SiamMask',
-#                 'experiments', 'siammask_sharp'))
-# sys.path.append(join(project_name, 'SiamMask', 'models'))
-# sys.path.append(join(project_name, 'SiamMask'))
-
 exp_path = join(project_name, 'SiamMask/experiments/siammask_sharp')
 pretrained_path1 = join(exp_path, 'SiamMask_DAVIS.pth')
 
-
 print(sys.path)
 
 torch.set_grad_enabled(False)
@@ -58,8 +52,7 @@ object_y = 0
 object_width = 0
 object_height = 0
 in_fps = 24
-original_frame_list = []
-mask_list = []
+
 
 parser = argparse.ArgumentParser()
 # parser.add_argument('--opt', default='configs/object_removal.yaml',
@@ -151,7 +144,7 @@ def getBoundaries(mask):
     return x1, y1, (x2-x1), (y2-y1)
 
 
-def track_and_mask(vid, original_frame, masked_frame):
+def track_and_mask(vid, masked_frame, original_list, mask_list):
     x, y, w, h = getBoundaries(masked_frame)
     f = 0
 
@@ -183,7 +176,7 @@ def track_and_mask(vid, original_frame, masked_frame):
                 # track
                 state = siamese_track(
                     state, frame, mask_enable=True, refine_enable=True, device=device)
-                original_frame_list.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+                original_list.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                 location = state['ploygon'].flatten()
                 mask = state['mask'] > state['p'].seg_thr
                 frame[:, :, 2] = (mask > 0) * 255 + \
@@ -206,17 +199,19 @@ def track_and_mask(vid, original_frame, masked_frame):
 
     print('Original Frame Count: ',len(original_frame_list))
     print('Mask Frame Count: ',len(mask_list))
-    return dt_string+"_output.avi"
+    return {out_video_inpaint:dt_string+"_output.avi",
+            original_frame_list: original_list,
+            mask_list: mask_list}
 
 
-def inpaint_video():
+def inpaint_video(original_frame_list, mask_list):
     args.out_fps = in_fps
     video_inpainting(args, original_frame_list, mask_list)
-    
-    original_frame_list.clear()
-    mask_list.clear()
-    
-    return dt_string+"_result.mp4"
+    original_frame_list = []
+    mask_list = []
+    return {out_video_inpaint:dt_string+"_result.mp4",
+            original_frame_list: original_frame_list,
+            mask_list: mask_list}
 
 
 def get_first_frame(video):
@@ -253,6 +248,8 @@ def getStartEndPoints(mask):
 
 
 with gr.Blocks() as demo:
+    original_frame_list = gr.State([])
+    mask_list = gr.State([])
     with gr.Row():
         with gr.Column(scale=2):
             with gr.Row():
@@ -277,8 +274,8 @@ with gr.Blocks() as demo:
     approve_mask.click(lambda x: [x['image'], x['mask']], first_frame, [
                        original_image, masked_image])
     track_mask.click(fn=track_and_mask, inputs=[
-                     in_video, original_image, masked_image], outputs=[out_video])
-    inpaint.click(fn=inpaint_video, outputs=[out_video_inpaint])
+                     in_video, masked_image, original_frame_list, mask_list], outputs=[out_video, original_frame_list, mask_list])
+    inpaint.click(fn=inpaint_video, outputs=[out_video_inpaint, original_frame_list, mask_list])
 
 
 demo.launch(debug=True)