RAM and VRAM usage reduce

app.py

@@ -341,7 +341,6 @@ def inpaint_video(video_state, interactive_state, mask_dropdown):
         operation_log = [("Error! You are trying to inpaint without masks input. Please track the selected mask first, and then press inpaint. If VRAM exceeded, please use the resize ratio to scaling down the image size.","Error"), ("","")]
         inpainted_frames = video_state["origin_images"]
         video_output = generate_video_from_frames(inpainted_frames, output_path="./result/inpaint/{}".format(video_state["video_name"]), fps=fps) # import video_input to name the output video
-    
     return video_output, operation_log
 
 
@@ -423,7 +422,7 @@ SAM_checkpoint = download_checkpoint(sam_checkpoint_url, folder, sam_checkpoint)
 xmem_checkpoint = download_checkpoint(xmem_checkpoint_url, folder, xmem_checkpoint)
 e2fgvi_checkpoint = download_checkpoint_from_google_drive(e2fgvi_checkpoint_id, folder, e2fgvi_checkpoint)
 # args.port = 12213
-# args.device = "cuda:1"
+# args.device = "cuda:8"
 # args.mask_save = True
 
 # initialize sam, xmem, e2fgvi models
@@ -432,7 +431,7 @@ model = TrackingAnything(SAM_checkpoint, xmem_checkpoint, e2fgvi_checkpoint,args
 
 title = """<p><h1 align="center">Track-Anything</h1></p>
     """
-description = """<p>Gradio demo for Track Anything, a flexible and interactive tool for video object tracking, segmentation, and inpainting. I To use it, simply upload your video, or click one of the examples to load them. Code: <a href="https://github.com/gaomingqi/Track-Anything">https://github.com/gaomingqi/Track-Anything</a> <a href="https://huggingface.co/spaces/watchtowerss/Track-Anything?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a></p>"""
+description = """<p>Gradio demo for Track Anything, a flexible and interactive tool for video object tracking, segmentation, and inpainting. To use it, simply upload your video, or click one of the examples to load them. Code: <a href="https://github.com/gaomingqi/Track-Anything">Track-Anything</a> <a href="https://huggingface.co/spaces/watchtowerss/Track-Anything?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a> If you stuck in unknown errors, please feel free to watch the Tutorial video.</p>"""
 
 
 with gr.Blocks() as iface:
@@ -450,7 +449,7 @@ with gr.Blocks() as iface:
             "masks": []
         },
         "track_end_number": None,
-        "resize_ratio": 1
+        "resize_ratio": 0.6
     }
     )
 
@@ -470,48 +469,78 @@ with gr.Blocks() as iface:
     gr.Markdown(title)
     gr.Markdown(description)
     with gr.Row():
-
-        # for user video input
         with gr.Column():
-            with gr.Row(scale=0.4):
-                video_input = gr.Video(autosize=True)
-                with gr.Column():
-                    video_info = gr.Textbox(label="Video Info")
-                    resize_info = gr.Textbox(value="If you want to use the inpaint function, it is best to git clone the repo and use a machine with more VRAM locally. \
-                                            Alternatively, you can use the resize ratio slider to scale down the original image to around 360P resolution for faster processing.", label="Tips for running this demo.")
-                    resize_ratio_slider = gr.Slider(minimum=0.02, maximum=1, step=0.02, value=1, label="Resize ratio", visible=True)
-          
-
-            with gr.Row():
-                # put the template frame under the radio button
+            with gr.Tab("Test"):
+        # for user video input
                 with gr.Column():
-                    # extract frames
-                    with gr.Column():
-                        extract_frames_button = gr.Button(value="Get video info", interactive=True, variant="primary") 
+                    with gr.Row(scale=0.4):
+                        video_input = gr.Video(autosize=True)
+                        with gr.Column():
+                            video_info = gr.Textbox(label="Video Info")
+                            resize_info = gr.Textbox(value="If you want to use the inpaint function, it is best to git clone the repo and use a machine with more VRAM locally. \
+                                                    Alternatively, you can use the resize ratio slider to scale down the original image to around 360P resolution for faster processing.", label="Tips for running this demo.")
+                            resize_ratio_slider = gr.Slider(minimum=0.02, maximum=1, step=0.02, value=0.6, label="Resize ratio", visible=True)
+                
 
-                     # click points settins, negative or positive, mode continuous or single
                     with gr.Row():
-                        with gr.Row():
-                            point_prompt = gr.Radio(
-                                choices=["Positive",  "Negative"],
-                                value="Positive",
-                                label="Point Prompt",
-                                interactive=True,
-                                visible=False)
-                            remove_mask_button = gr.Button(value="Remove mask", interactive=True, visible=False) 
-                            clear_button_click = gr.Button(value="Clear Clicks", interactive=True, visible=False).style(height=160)
-                            Add_mask_button = gr.Button(value="Add mask", interactive=True, visible=False)
-                    template_frame = gr.Image(type="pil",interactive=True, elem_id="template_frame", visible=False).style(height=360)
-                    image_selection_slider = gr.Slider(minimum=1, maximum=100, step=1, value=1, label="Image Selection", visible=False)
-                    track_pause_number_slider = gr.Slider(minimum=1, maximum=100, step=1, value=1, label="Track end frames", visible=False)
-            
+                        # put the template frame under the radio button
+                        with gr.Column():
+                            # extract frames
+                            with gr.Column():
+                                extract_frames_button = gr.Button(value="Get video info", interactive=True, variant="primary") 
+
+                            # click points settins, negative or positive, mode continuous or single
+                            with gr.Row():
+                                with gr.Row():
+                                    point_prompt = gr.Radio(
+                                        choices=["Positive",  "Negative"],
+                                        value="Positive",
+                                        label="Point Prompt",
+                                        interactive=True,
+                                        visible=False)
+                                    remove_mask_button = gr.Button(value="Remove mask", interactive=True, visible=False) 
+                                    clear_button_click = gr.Button(value="Clear Clicks", interactive=True, visible=False).style(height=160)
+                                    Add_mask_button = gr.Button(value="Add mask", interactive=True, visible=False)
+                            template_frame = gr.Image(type="pil",interactive=True, elem_id="template_frame", visible=False).style(height=360)
+                            image_selection_slider = gr.Slider(minimum=1, maximum=100, step=1, value=1, label="Image Selection", visible=False)
+                            track_pause_number_slider = gr.Slider(minimum=1, maximum=100, step=1, value=1, label="Track end frames", visible=False)
+                    
+                        with gr.Column():
+                            run_status = gr.HighlightedText(value=[("Run","Error"),("Status","Normal")], visible=True)
+                            mask_dropdown = gr.Dropdown(multiselect=True, value=[], label="Mask selection", info=".", visible=False)
+                            video_output = gr.Video(autosize=True, visible=False).style(height=360)
+                            with gr.Row():
+                                tracking_video_predict_button = gr.Button(value="Tracking", visible=False)
+                                inpaint_video_predict_button = gr.Button(value="Inpaint", visible=False)
+                # set example
+                gr.Markdown("##  Examples")
+                gr.Examples(
+                    examples=[os.path.join(os.path.dirname(__file__), "./test_sample/", test_sample) for test_sample in ["test-sample8.mp4","test-sample4.mp4", \
+                                                                                                                        "test-sample2.mp4","test-sample13.mp4"]],
+                    fn=run_example,
+                    inputs=[
+                        video_input
+                    ],
+                    outputs=[video_input],
+                    # cache_examples=True,
+    ) 
+                
+            with gr.Tab("Tutorial"):
                 with gr.Column():
-                    run_status = gr.HighlightedText(value=[("Text","Error"),("to be","Label 2"),("highlighted","Label 3")], visible=True)
-                    mask_dropdown = gr.Dropdown(multiselect=True, value=[], label="Mask selection", info=".", visible=False)
-                    video_output = gr.Video(autosize=True, visible=False).style(height=360)
-                    with gr.Row():
-                        tracking_video_predict_button = gr.Button(value="Tracking", visible=False)
-                        inpaint_video_predict_button = gr.Button(value="Inpaint", visible=False)
+                    with gr.Row(scale=0.4):
+                        video_demo_operation = gr.Video(autosize=True) 
+                
+                # set example
+                gr.Markdown("## Operation tutorial video")
+                gr.Examples(
+                    examples=[os.path.join(os.path.dirname(__file__), "./test_sample/", test_sample) for test_sample in ["huggingface_demo_operation.mp4"]],
+                    fn=run_example,
+                    inputs=[
+                        video_demo_operation
+                    ],
+                    outputs=[video_demo_operation],
+                    # cache_examples=True,
+                ) 
 
     # first step: get the video information 
     extract_frames_button.click(
@@ -601,7 +630,7 @@ with gr.Blocks() as iface:
             "masks": []
         },
         "track_end_number": 0,
-        "resize_ratio": 1
+        "resize_ratio": 0.6
         },
         [[],[]],
         None,
@@ -609,7 +638,7 @@ with gr.Blocks() as iface:
         gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), \
         gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), \
         gr.update(visible=False), gr.update(visible=False), gr.update(visible=False, value=[]), gr.update(visible=False), \
-        gr.update(visible=False), gr.update(visible=False)
+        gr.update(visible=False), gr.update(visible=True)
                         
         ),
         [],
@@ -631,18 +660,6 @@ with gr.Blocks() as iface:
         inputs = [video_state, click_state,],
         outputs = [template_frame,click_state, run_status],
     )
-    # set example
-    gr.Markdown("##  Examples")
-    gr.Examples(
-        examples=[os.path.join(os.path.dirname(__file__), "./test_sample/", test_sample) for test_sample in ["test-sample8.mp4","test-sample4.mp4", \
-                                                                                                             "test-sample2.mp4","test-sample13.mp4"]],
-        fn=run_example,
-        inputs=[
-            video_input
-        ],
-        outputs=[video_input],
-        # cache_examples=True,
-    ) 
 iface.queue(concurrency_count=1)
 # iface.launch(debug=True, enable_queue=True, server_port=args.port, server_name="0.0.0.0")
 iface.launch(debug=True, enable_queue=True)

inpainter/base_inpainter.py

@@ -9,9 +9,9 @@ import numpy as np
 from tqdm import tqdm
 from inpainter.util.tensor_util import resize_frames, resize_masks
 
-def read_image_from_userfolder(image_path):
+def read_image_from_split(videp_split_path):
     # if type:
-    image = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
+    image = np.asarray([np.asarray(Image.open(path)) for path in videp_split_path])
     # else:
         # image = cv2.cvtColor(cv2.imread("/tmp/{}/paintedimages/{}/{:08d}.png".format(username, video_state["video_name"], index+ ".png")), cv2.COLOR_BGR2RGB)
     return image
@@ -56,29 +56,30 @@ class BaseInpainter:
                         break
                     ref_index.append(i)
         return ref_index
-
-    def inpaint(self, frames_path, masks, dilate_radius=15, ratio=1):
+    
+    def inpaint_efficient(self, frames, masks, num_tcb, num_tca, dilate_radius=15, ratio=1):
         """
+        Perform Inpainting for video subsets
         frames: numpy array, T, H, W, 3
         masks: numpy array, T, H, W
+        num_tcb: constant, number of temporal context before, frames
+        num_tca: constant, number of temporal context after, frames
         dilate_radius: radius when applying dilation on masks
         ratio: down-sample ratio
 
         Output:
         inpainted_frames: numpy array, T, H, W, 3
         """
-        frames = []
-        for file in frames_path:
-            frames.append(read_image_from_userfolder(file))
-        frames = np.asarray(frames)
-
         assert frames.shape[:3] == masks.shape, 'different size between frames and masks'
         assert ratio > 0 and ratio <= 1, 'ratio must in (0, 1]'
+        
+        # --------------------
+        # pre-processing
+        # --------------------
         masks = masks.copy()
         masks = np.clip(masks, 0, 1)
         kernel = cv2.getStructuringElement(2, (dilate_radius, dilate_radius))
         masks = np.stack([cv2.dilate(mask, kernel) for mask in masks], 0)
-
         T, H, W = masks.shape
         masks = np.expand_dims(masks, axis=3)    # expand to T, H, W, 1
         # size: (w, h)
@@ -96,14 +97,37 @@ class BaseInpainter:
             frames = resize_frames(frames, tuple(size))          # T, H, W, 3
         # frames and binary_masks are numpy arrays
         h, w = frames.shape[1:3]
-        video_length = T
-
+        video_length = T - (num_tca + num_tcb)  # real video length
         # convert to tensor
         imgs = (torch.from_numpy(frames).permute(0, 3, 1, 2).contiguous().unsqueeze(0).float().div(255)) * 2 - 1
         masks = torch.from_numpy(binary_masks).permute(0, 3, 1, 2).contiguous().unsqueeze(0)
-
         imgs, masks = imgs.to(self.device), masks.to(self.device)
         comp_frames = [None] * video_length
+        tcb_imgs = None
+        tca_imgs = None
+        tcb_masks = None
+        tca_masks = None
+        # --------------------
+        # end of pre-processing
+        # --------------------
+
+        # separate tc frames/masks from imgs and masks
+        if num_tcb > 0:
+            tcb_imgs = imgs[:, :num_tcb]
+            tcb_masks = masks[:, :num_tcb]
+            tcb_binary = binary_masks[:num_tcb]
+        if num_tca > 0:
+            tca_imgs = imgs[:, -num_tca:]
+            tca_masks = masks[:, -num_tca:]
+            tca_binary = binary_masks[-num_tca:]
+            end_idx = -num_tca
+        else:
+            end_idx = T
+
+        imgs = imgs[:, num_tcb:end_idx]
+        masks = masks[:, num_tcb:end_idx]
+        binary_masks = binary_masks[num_tcb:end_idx]	# only neighbor area are involved
+        frames = frames[num_tcb:end_idx]				# only neighbor area are involved
 
         for f in tqdm(range(0, video_length, self.neighbor_stride), desc='Inpainting image'):
             neighbor_ids = [
@@ -111,8 +135,24 @@ class BaseInpainter:
                                 min(video_length, f + self.neighbor_stride + 1))
             ]
             ref_ids = self.get_ref_index(f, neighbor_ids, video_length)
-            selected_imgs = imgs[:1, neighbor_ids + ref_ids, :, :, :]
-            selected_masks = masks[:1, neighbor_ids + ref_ids, :, :, :]
+
+            # selected_imgs = imgs[:1, neighbor_ids + ref_ids, :, :, :]
+            # selected_masks = masks[:1, neighbor_ids + ref_ids, :, :, :]
+            
+            selected_imgs = imgs[:, neighbor_ids]
+            selected_masks = masks[:, neighbor_ids]
+            # pad before
+            if tcb_imgs is not None:
+                selected_imgs = torch.concat([selected_imgs, tcb_imgs], dim=1)
+                selected_masks = torch.concat([selected_masks, tcb_masks], dim=1)
+            # integrate ref frames
+            selected_imgs = torch.concat([selected_imgs, imgs[:, ref_ids]], dim=1)
+            selected_masks = torch.concat([selected_masks, masks[:, ref_ids]], dim=1)
+            # pad after
+            if tca_imgs is not None:
+                selected_imgs = torch.concat([selected_imgs, tca_imgs], dim=1)
+                selected_masks = torch.concat([selected_masks, tca_masks], dim=1)
+
             with torch.no_grad():
                 masked_imgs = selected_imgs * (1 - selected_masks)
                 mod_size_h = 60
@@ -138,10 +178,75 @@ class BaseInpainter:
                     else:
                         comp_frames[idx] = comp_frames[idx].astype(
                             np.float32) * 0.5 + img.astype(np.float32) * 0.5
-        
+            torch.cuda.empty_cache()
         inpainted_frames = np.stack(comp_frames, 0)
         return inpainted_frames.astype(np.uint8)
 
+    def inpaint(self, frames_path, masks, dilate_radius=15, ratio=1):
+        """
+		Perform Inpainting for video subsets
+		frames: numpy array, T, H, W, 3
+		masks: numpy array, T, H, W
+		dilate_radius: radius when applying dilation on masks
+		ratio: down-sample ratio
+
+		Output:
+		inpainted_frames: numpy array, T, H, W, 3
+		"""
+        # assert frames.shape[:3] == masks.shape, 'different size between frames and masks'
+        assert ratio > 0 and ratio <= 1, 'ratio must in (0, 1]'
+        
+        # set interval
+        interval = 45
+        context_range = 10  # for each split, consider its temporal context [-context_range] frames and [context_range] frames
+        # split frames into subsets
+        video_length = len(frames_path)
+        num_splits = video_length // interval
+        id_splits = [[i*interval, (i+1)*interval] for i in range(num_splits)]  # id splits
+        # if remaining split > interval/2, add a new split, else, append to the last split
+        if video_length - id_splits[-1][-1] > interval / 2:
+            id_splits.append([num_splits*interval, video_length])
+        else:
+            id_splits[-1][-1] = video_length
+
+        # perform inpainting for each split
+        inpainted_splits = []
+        for id_split in id_splits:
+            video_split_path = frames_path[id_split[0]:id_split[1]]
+            video_split = read_image_from_split(video_split_path)
+            mask_split = masks[id_split[0]:id_split[1]]
+
+            # | id_before | ----- | id_split[0] | ----- | id_split[1] | ----- | id_after |
+            # add temporal context 
+            id_before = max(0, id_split[0] - self.step * context_range)
+            try:
+                tcb_frames = np.stack([np.array(Image.open(frames_path[idb])) for idb in range(id_before, id_split[0]-self.step, self.step)], 0)
+                tcb_masks = np.stack([masks[idb] for idb in range(id_before, id_split[0]-self.step, self.step)], 0)
+                num_tcb = len(tcb_frames)
+            except:
+                num_tcb = 0
+            id_after = min(video_length, id_split[1] + self.step * context_range)
+            try:
+                tca_frames = np.stack([np.array(Image.open(frames_path[ida])) for ida in range(id_split[1]+self.step, id_after, self.step)], 0)
+                tca_masks = np.stack([masks[ida] for ida in range(id_split[1]+self.step, id_after, self.step)], 0)
+                num_tca = len(tca_frames)
+            except:
+                num_tca = 0
+
+            # concatenate temporal context frames/masks with input frames/masks (for parallel pre-processing)
+            if num_tcb > 0:
+                video_split = np.concatenate([tcb_frames, video_split], 0)
+                mask_split = np.concatenate([tcb_masks, mask_split], 0)
+            if num_tca > 0:
+                video_split = np.concatenate([video_split, tca_frames], 0)
+                mask_split = np.concatenate([mask_split, tca_masks], 0)
+            
+            # inpaint each split
+            inpainted_splits.append(self.inpaint_efficient(video_split, mask_split, num_tcb, num_tca, dilate_radius, ratio))
+
+        inpainted_frames = np.concatenate(inpainted_splits, 0)
+        return inpainted_frames.astype(np.uint8) 
+
 if __name__ == '__main__':
 
     frame_path = glob.glob(os.path.join('/ssd1/gaomingqi/datasets/davis/JPEGImages/480p/parkour', '*.jpg'))
@@ -179,4 +284,4 @@ if __name__ == '__main__':
     # save
     for ti, inpainted_frame in enumerate(inpainted_frames):
         frame = Image.fromarray(inpainted_frame).convert('RGB')
-        frame.save(os.path.join(save_path, f'{ti:05d}.jpg'))
+        frame.save(os.path.join(save_path, f'{ti:05d}.jpg'))

test_sample/test-sample13.mp4

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bf112202beb75ecf7d04b27758f1f3eedfc218dac5d5dad0b72a07dd2db0f423
-size 59659465
+oid sha256:54f24a0aaae482aff7ff3555256f60ad1931d478dc5694fb37624cac85479eee
+size 2528426

test_sample/test-sample8.mp4

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:01d255ef82222d950d2cfae904d82cc20c752577016f0325b21788fb9b458bb9
-size 11979994
+oid sha256:2414d24cc1ddfe1619c17e9876a7c3ed0f1f37da234c63c08af2cecbbb16c1ed
+size 8714250