Merge branch 'master' of https://github.com/gaomingqi/VOS-Anything

app.py

@@ -18,6 +18,7 @@ import torch
 import concurrent.futures
 import queue
 
+# download checkpoints
 def download_checkpoint(url, folder, filename):
     os.makedirs(folder, exist_ok=True)
     filepath = os.path.join(folder, filename)
@@ -34,16 +35,6 @@ def download_checkpoint(url, folder, filename):
 
     return filepath
 
-def pause_video(play_state):
-    print("user pause_video")
-    play_state.append(time.time())
-    return play_state
-
-def play_video(play_state):
-    print("user play_video")
-    play_state.append(time.time())
-    return play_state
-
 # convert points input to prompt state
 def get_prompt(click_state, click_input):
     inputs = json.loads(click_input)
@@ -61,8 +52,9 @@ def get_prompt(click_state, click_input):
         "multimask_output":"True",
     }
     return prompt
-    
+
-def get_frames_from_video(video_input, play_state):
+# extract frames from upload video
+def get_frames_from_video(video_input, video_state):
     """
     Args:
         video_path:str
@@ -71,10 +63,6 @@ def get_frames_from_video(video_input, play_state):
         [[0:nearest_frame], [nearest_frame:], nearest_frame]
     """
     video_path = video_input
-    try:
-        timestamp = play_state[1] - play_state[0]
-    except:
-        timestamp = 0.1
     frames = []
     try:
         cap = cv2.VideoCapture(video_path)
@@ -88,19 +76,97 @@ def get_frames_from_video(video_input, play_state):
     except (OSError, TypeError, ValueError, KeyError, SyntaxError) as e:
         print("read_frame_source:{} error. {}\n".format(video_path, str(e)))
 
-    # for index, frame in enumerate(frames):
+    # initialize video_state
-        # frames[index] = np.asarray(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
+    video_state = {
+        "video_name": os.path.split(video_path)[-1],
+        "origin_images": frames,
+        "painted_images": frames.copy(),
+        "masks": [None]*len(frames),
+        "logits": [None]*len(frames),
+        "select_frame_number": 0,
+        "fps": 30
+        }
+    return video_state, gr.update(visible=True, maximum=len(frames), value=1)
+
+# get the select frame from gradio slider
+def select_template(image_selection_slider, video_state):
+
+    # images = video_state[1]
+    image_selection_slider -= 1
+    video_state["select_frame_number"] = image_selection_slider
+
+    # once select a new template frame, set the image in sam
+
+    model.samcontroler.sam_controler.reset_image()
+    model.samcontroler.sam_controler.set_image(video_state["origin_images"][image_selection_slider])
+
+
+    return video_state["painted_images"][image_selection_slider], video_state
+
+# use sam to get the mask
+def sam_refine(video_state, point_prompt, click_state, interactive_state, evt:gr.SelectData):
+    """
+    Args:
+        template_frame: PIL.Image
+        point_prompt: flag for positive or negative button click
+        click_state: [[points], [labels]]
+    """
+    if point_prompt == "Positive":
+        coordinate = "[[{},{},1]]".format(evt.index[0], evt.index[1])
+        interactive_state["positive_click_times"] += 1
+    else:
+        coordinate = "[[{},{},0]]".format(evt.index[0], evt.index[1])
+        interactive_state["negative_click_times"] += 1
     
-    key_frame_index = int(timestamp * fps)
+    # prompt for sam model
-    nearest_frame = frames[key_frame_index]
+    prompt = get_prompt(click_state=click_state, click_input=coordinate)
-    frames_split = [frames[:key_frame_index], frames[key_frame_index:], nearest_frame]
-    # output_path='./seperate.mp4'
-    # torchvision.io.write_video(output_path, frames[1], fps=fps, video_codec="libx264")
 
-    # set image in sam when select the template frame
+    mask, logit, painted_image = model.first_frame_click( 
-    model.samcontroler.sam_controler.set_image(nearest_frame)
+                                                      image=video_state["origin_images"][video_state["select_frame_number"]], 
-    return frames_split, nearest_frame, nearest_frame, fps
+                                                      points=np.array(prompt["input_point"]),
+                                                      labels=np.array(prompt["input_label"]),
+                                                      multimask=prompt["multimask_output"],
+                                                      )
+    video_state["masks"][video_state["select_frame_number"]] = mask
+    video_state["logits"][video_state["select_frame_number"]] = logit
+    video_state["painted_images"][video_state["select_frame_number"]] = painted_image
 
+    return painted_image, video_state, interactive_state
+
+# tracking vos
+def vos_tracking_video(video_state, interactive_state):
+    model.xmem.clear_memory()
+    following_frames = video_state["origin_images"][video_state["select_frame_number"]:]
+    template_mask = video_state["masks"][video_state["select_frame_number"]]
+    fps = video_state["fps"]
+    masks, logits, painted_images = model.generator(images=following_frames, template_mask=template_mask)
+
+    video_state["masks"][video_state["select_frame_number"]:] = masks
+    video_state["logits"][video_state["select_frame_number"]:] = logits
+    video_state["painted_images"][video_state["select_frame_number"]:] = painted_images
+
+    video_output = generate_video_from_frames(video_state["painted_images"], output_path="./result/{}".format(video_state["video_name"]), fps=fps) # import video_input to name the output video
+    interactive_state["inference_times"] += 1
+    
+    print("For generating this tracking result, inference times: {}, click times: {}, positive: {}, negative: {}".format(interactive_state["inference_times"], 
+                                                                                                                                           interactive_state["positive_click_times"]+interactive_state["negative_click_times"],
+                                                                                                                                           interactive_state["positive_click_times"],
+                                                                                                                                        interactive_state["negative_click_times"]))
+    
+    #### shanggao code for mask save
+    if interactive_state["mask_save"]:
+        if not os.path.exists('./result/mask/{}'.format(video_state["video_name"].split('.')[0])):
+            os.makedirs('./result/mask/{}'.format(video_state["video_name"].split('.')[0]))
+        i = 0
+        print("save mask")
+        for mask in video_state["masks"]:
+            np.save(os.path.join('./result/mask/{}'.format(video_state["video_name"].split('.')[0]), '{:05d}.npy'.format(i)), mask)
+            i+=1
+        # save_mask(video_state["masks"], video_state["video_name"])
+    #### shanggao code for mask save
+    return video_output, video_state, interactive_state
+
+# generate video after vos inference
 def generate_video_from_frames(frames, output_path, fps=30):
     """
     Generates a video from a list of frames.
@@ -110,104 +176,12 @@ def generate_video_from_frames(frames, output_path, fps=30):
         output_path (str): The path to save the generated video.
         fps (int, optional): The frame rate of the output video. Defaults to 30.
     """
-    # height, width, layers = frames[0].shape
-    # fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    # video = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
-    
-    # for frame in frames:
-    #     video.write(frame)
-    
-    # video.release()
     frames = torch.from_numpy(np.asarray(frames))
-    output_path='./output.mp4'
+    if not os.path.exists(os.path.dirname(output_path)):
+        os.makedirs(os.path.dirname(output_path))
     torchvision.io.write_video(output_path, frames, fps=fps, video_codec="libx264")
     return output_path
 
-def model_reset():
-    model.xmem.clear_memory()
-    return None
-
-def sam_refine(origin_frame, point_prompt, click_state, logit, evt:gr.SelectData):
-    """
-    Args:
-        template_frame: PIL.Image
-        point_prompt: flag for positive or negative button click
-        click_state: [[points], [labels]]
-    """
-    if point_prompt == "Positive":
-        coordinate = "[[{},{},1]]".format(evt.index[0], evt.index[1])
-    else:
-        coordinate = "[[{},{},0]]".format(evt.index[0], evt.index[1])
-    
-    # prompt for sam model
-    prompt = get_prompt(click_state=click_state, click_input=coordinate)
-
-    # default value
-    # points = np.array([[evt.index[0],evt.index[1]]])
-    # labels= np.array([1])
-    if len(logit)==0:
-        logit = None
-    
-    mask, logit, painted_image = model.first_frame_click( 
-                                                      image=origin_frame, 
-                                                      points=np.array(prompt["input_point"]),
-                                                      labels=np.array(prompt["input_label"]),
-                                                      multimask=prompt["multimask_output"],
-                                                      )
-    return painted_image, click_state, logit, mask
-
-
-
-def vos_tracking_video(video_state, template_mask,fps):
-
-    masks, logits, painted_images = model.generator(images=video_state[1], template_mask=template_mask)
-    video_output = generate_video_from_frames(painted_images, output_path="./output.mp4", fps=fps)
-    # image_selection_slider = gr.Slider(minimum=1, maximum=len(video_state[1]), value=1, label="Image Selection", interactive=True)
-    return video_output, painted_images, masks, logits
-
-def vos_tracking_image(image_selection_slider, painted_images):
-
-    # images = video_state[1]
-    percentage = image_selection_slider / 100
-    select_frame_num = int(percentage * len(painted_images))
-    return painted_images[select_frame_num], select_frame_num
-
-def interactive_correction(video_state, point_prompt, click_state, select_correction_frame, evt: gr.SelectData):
-    """
-    Args:
-        template_frame: PIL.Image
-        point_prompt: flag for positive or negative button click
-        click_state: [[points], [labels]]
-    """
-    refine_image = video_state[1][select_correction_frame]
-    if point_prompt == "Positive":
-        coordinate = "[[{},{},1]]".format(evt.index[0], evt.index[1])
-    else:
-        coordinate = "[[{},{},0]]".format(evt.index[0], evt.index[1])
-    
-    # prompt for sam model
-    prompt = get_prompt(click_state=click_state, click_input=coordinate)
-    model.samcontroler.seg_again(refine_image)
-    corrected_mask, corrected_logit, corrected_painted_image = model.first_frame_click( 
-                                                      image=refine_image, 
-                                                      points=np.array(prompt["input_point"]),
-                                                      labels=np.array(prompt["input_label"]),
-                                                      multimask=prompt["multimask_output"],
-                                                      )
-    return corrected_painted_image, [corrected_mask, corrected_logit, corrected_painted_image]
-
-def correct_track(video_state, select_correction_frame, corrected_state, masks, logits, painted_images, fps):
-    model.xmem.clear_memory()
-    # inference the following images
-    following_images = video_state[1][select_correction_frame:]
-    corrected_masks, corrected_logits, corrected_painted_images = model.generator(images=following_images, template_mask=corrected_state[0])
-    masks = masks[:select_correction_frame] + corrected_masks
-    logits = logits[:select_correction_frame] + corrected_logits
-    painted_images = painted_images[:select_correction_frame] + corrected_painted_images
-    video_output = generate_video_from_frames(painted_images, output_path="./output.mp4", fps=fps)
-
-    return video_output, painted_images, logits, masks 
-
 # check and download checkpoints if needed
 SAM_checkpoint = "sam_vit_h_4b8939.pth" 
 sam_checkpoint_url = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
@@ -219,8 +193,9 @@ xmem_checkpoint = download_checkpoint(xmem_checkpoint_url, folder, xmem_checkpoi
 
 # args, defined in track_anything.py
 args = parse_augment()
-args.port = 12219
+args.port = 12212
-args.device = "cuda:5"
+args.device = "cuda:4"
+args.mask_save = True
 
 model = TrackingAnything(SAM_checkpoint, xmem_checkpoint, args)
 
@@ -228,36 +203,40 @@ with gr.Blocks() as iface:
     """
         state for 
     """
-    state = gr.State([])
-    play_state = gr.State([])
-    video_state = gr.State([[],[],[]])
     click_state = gr.State([[],[]])
-    logits = gr.State([])
+    interactive_state = gr.State({
-    masks = gr.State([])
+        "inference_times": 0,
-    painted_images = gr.State([])
+        "negative_click_times" : 0,
-    origin_image = gr.State(None)
+        "positive_click_times": 0,
-    template_mask = gr.State(None)
+        "mask_save": args.mask_save
-    select_correction_frame = gr.State(None)
+    })
-    corrected_state = gr.State([[],[],[]])
+    video_state = gr.State(
-    fps = gr.State([])
+        {
-    # queue value for image refresh, origin image, mask, logits, painted image
+        "video_name": "",
-
+        "origin_images": None,
-
+        "painted_images": None,
+        "masks": None,
+        "logits": None,
+        "select_frame_number": 0,
+        "fps": 30
+        }
+    )
 
     with gr.Row():
 
         # for user video input
         with gr.Column(scale=1.0):
-            video_input = gr.Video().style(height=720)
+            video_input = gr.Video().style(height=360)
 
-            # listen to the user action for play and pause input video
-            video_input.play(fn=play_video, inputs=play_state, outputs=play_state, scroll_to_output=True, show_progress=True)
-            video_input.pause(fn=pause_video, inputs=play_state, outputs=play_state)
           
 
             with gr.Row(scale=1):
                 # put the template frame under the radio button
                 with gr.Column(scale=0.5):
+                    # 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(scale=0.5):
@@ -275,75 +254,99 @@ with gr.Blocks() as iface:
                             clear_button_clike = gr.Button(value="Clear Clicks", interactive=True).style(height=160)
                             clear_button_image = gr.Button(value="Clear Image", interactive=True)
                     template_frame = gr.Image(type="pil",interactive=True, elem_id="template_frame").style(height=360)
-                    with gr.Column():
+                    image_selection_slider = gr.Slider(minimum=1, maximum=100, step=1, value=1, label="Image Selection", invisible=False)
-                        template_select_button = gr.Button(value="Template select", interactive=True, variant="primary")
+
-                    
+
                    
             
                 with gr.Column(scale=0.5):
-
-
-                    # for intermedia result check and correction
-                    # intermedia_image = gr.Image(type="pil", interactive=True, elem_id="intermedia_frame").style(height=360)
                     video_output = gr.Video().style(height=360)
                     tracking_video_predict_button = gr.Button(value="Tracking")
 
-                    image_output = gr.Image(type="pil", interactive=True, elem_id="image_output").style(height=360)
+    # first step: get the video information 
-                    image_selection_slider = gr.Slider(minimum=0, maximum=100, step=0.1, value=0, label="Image Selection", interactive=True)
+    extract_frames_button.click(
-                    correct_track_button = gr.Button(value="Interactive Correction")
+        fn=get_frames_from_video,
+        inputs=[
+            video_input, video_state
+        ],
+        outputs=[video_state, image_selection_slider],
+    )   
+
+    # second step: select images from slider
+    image_selection_slider.release(fn=select_template, 
+                                   inputs=[image_selection_slider, video_state], 
+                                   outputs=[template_frame, video_state], api_name="select_image")
+    
 
     template_frame.select(
         fn=sam_refine,
-        inputs=[
+        inputs=[video_state, point_prompt, click_state, interactive_state],
-            origin_image, point_prompt, click_state, logits
+        outputs=[template_frame, video_state, interactive_state]
-        ],
-        outputs=[
-            template_frame, click_state, logits, template_mask
-        ]
     )
-            
-    template_select_button.click(
-        fn=get_frames_from_video,
-        inputs=[
-            video_input, 
-            play_state
-        ],
-        outputs=[video_state, template_frame, origin_image, fps],
-    )   
 
     tracking_video_predict_button.click(
         fn=vos_tracking_video,
-        inputs=[video_state, template_mask, fps],
+        inputs=[video_state, interactive_state],
-        outputs=[video_output, painted_images, masks, logits]
+        outputs=[video_output, video_state, interactive_state]
-    )
-    image_selection_slider.release(fn=vos_tracking_image, 
-                                   inputs=[image_selection_slider, painted_images], outputs=[image_output, select_correction_frame], api_name="select_image")
-    # correction
-    image_output.select(
-        fn=interactive_correction,
-        inputs=[video_state, point_prompt, click_state, select_correction_frame],
-        outputs=[image_output, corrected_state]
-    )
-    correct_track_button.click(
-        fn=correct_track,
-        inputs=[video_state, select_correction_frame, corrected_state, masks, logits, painted_images, fps],
-        outputs=[video_output, painted_images, logits, masks ]
     )
+
     
     # clear input
     video_input.clear(
-        lambda: ([], [], [[], [], []], 
+        lambda: (
-                 None, "", "", "", "", "", "", "", [[],[]],
+        {
-                 None),
+        "origin_images": None,
+        "painted_images": None,
+        "masks": None,
+        "logits": None,
+        "select_frame_number": 0,
+        "fps": 30
+        },
+        {
+        "inference_times": 0,
+        "negative_click_times" : 0,
+        "positive_click_times": 0,
+        "mask_save": args.mask_save
+        },
+        [[],[]]
+                ),
         [],
-        [ state, play_state, video_state, 
+        [ 
-         template_frame, video_output, image_output, origin_image, template_mask, painted_images, masks, logits, click_state,
+            video_state,
-         select_correction_frame],
+            interactive_state,
+            click_state,
+        ],
         queue=False,
         show_progress=False
     )
     clear_button_image.click(
-        fn=model_reset
+        lambda: (
+        {
+        "origin_images": None,
+        "painted_images": None,
+        "masks": None,
+        "logits": None,
+        "select_frame_number": 0,
+        "fps": 30
+        },
+        { 
+        "inference_times": 0,
+        "negative_click_times" : 0,
+        "positive_click_times": 0,
+        "mask_save": args.mask_save
+        },
+        [[],[]]
+                ),
+        [],
+        [ 
+            video_state,
+            interactive_state,
+            click_state,
+        ],
+
+        queue=False,
+        show_progress=False
+
     )
     clear_button_clike.click(
        lambda: ([[],[]]),

track_anything.py

@@ -62,6 +62,7 @@ def parse_augment():
     parser.add_argument('--sam_model_type', type=str, default="vit_h")
     parser.add_argument('--port', type=int, default=6080, help="only useful when running gradio applications")  
     parser.add_argument('--debug', action="store_true")
+    parser.add_argument('--mask_save', default=True)
     args = parser.parse_args()
 
     if args.debug:

tracker/inference/memory_manager.py

@@ -182,7 +182,7 @@ class MemoryManager:
         if self.enable_long_term:
             # Do memory compressed if needed
             if self.work_mem.size >= self.max_work_elements:
-                print('remove memory')
+                # print('remove memory')
                 # Remove obsolete features if needed
                 if self.long_mem.size >= (self.max_long_elements-self.num_prototypes):
                     self.long_mem.remove_obsolete_features(self.max_long_elements-self.num_prototypes)
@@ -239,8 +239,8 @@ class MemoryManager:
 
         # add to long-term memory
         self.long_mem.add(prototype_key, prototype_value, prototype_shrinkage, selection=None, objects=None)
-        print(f'long memory size: {self.long_mem.size}')
+        # print(f'long memory size: {self.long_mem.size}')
-        print(f'work memory size: {self.work_mem.size}')
+        # print(f'work memory size: {self.work_mem.size}')
 
     def consolidation(self, candidate_key, candidate_shrinkage, candidate_selection, usage, candidate_value):
         # keys: 1*C*N