requirement fix

app.py

@@ -13,7 +13,13 @@ import requests
 import json
 import torchvision
 import torch 
+from tools.interact_tools import SamControler
+from tracker.base_tracker import BaseTracker
 from tools.painter import mask_painter
+try: 
+    from mmcv.cnn import ConvModule
+except:
+    os.system("mim install mmcv")
 
 # download checkpoints
 def download_checkpoint(url, folder, filename):
@@ -200,6 +206,7 @@ def show_mask(video_state, interactive_state, mask_dropdown):
 
 # tracking vos
 def vos_tracking_video(video_state, interactive_state, mask_dropdown):
+
     model.xmem.clear_memory()
     if interactive_state["track_end_number"]:
         following_frames = video_state["origin_images"][video_state["select_frame_number"]:interactive_state["track_end_number"]]
@@ -219,6 +226,8 @@ def vos_tracking_video(video_state, interactive_state, mask_dropdown):
         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)
+    # clear GPU memory
+    model.xmem.clear_memory()
 
     if interactive_state["track_end_number"]: 
         video_state["masks"][video_state["select_frame_number"]:interactive_state["track_end_number"]] = masks
@@ -258,6 +267,7 @@ def vos_tracking_video(video_state, interactive_state, mask_dropdown):
 
 # inpaint 
 def inpaint_video(video_state, interactive_state, mask_dropdown):
+
     frames = np.asarray(video_state["origin_images"])
     fps = video_state["fps"]
     inpaint_masks = np.asarray(video_state["masks"])
@@ -304,20 +314,33 @@ def generate_video_from_frames(frames, output_path, fps=30):
     torchvision.io.write_video(output_path, frames, fps=fps, video_codec="libx264")
     return output_path
 
+
+# args, defined in track_anything.py
+args = parse_augment()
+
 # 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"
+SAM_checkpoint_dict = {
+    'vit_h': "sam_vit_h_4b8939.pth",
+    'vit_l': "sam_vit_l_0b3195.pth", 
+    "vit_b": "sam_vit_b_01ec64.pth"
+}
+SAM_checkpoint_url_dict = {
+    'vit_h': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
+    'vit_l': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",
+    'vit_b': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
+}
+sam_checkpoint = SAM_checkpoint_dict[args.sam_model_type] 
+sam_checkpoint_url = SAM_checkpoint_url_dict[args.sam_model_type] 
 xmem_checkpoint = "XMem-s012.pth"
 xmem_checkpoint_url = "https://github.com/hkchengrex/XMem/releases/download/v1.0/XMem-s012.pth"
 e2fgvi_checkpoint = "E2FGVI-HQ-CVPR22.pth"
 e2fgvi_checkpoint_id = "10wGdKSUOie0XmCr8SQ2A2FeDe-mfn5w3"
 
+
 folder ="./checkpoints"
-SAM_checkpoint = download_checkpoint(sam_checkpoint_url, folder, SAM_checkpoint)
+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, defined in track_anything.py
-args = parse_augment()
 # args.port = 12315
 # args.device = "cuda:2"
 # args.mask_save = True
@@ -325,6 +348,12 @@ args = parse_augment()
 # initialize sam, xmem, e2fgvi models
 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>"""
+
+
 with gr.Blocks() as iface:
     """
         state for 
@@ -356,7 +385,8 @@ with gr.Blocks() as iface:
         "fps": 30
         }
     )
-
+    gr.Markdown(title)
+    gr.Markdown(description)
     with gr.Row():
 
         # for user video input
@@ -365,7 +395,7 @@ with gr.Blocks() as iface:
                 video_input = gr.Video(autosize=True)
                 with gr.Column():
                     video_info = gr.Textbox()
-                    video_info = gr.Textbox(value="Due to server restrictions, please upload a video that is no longer than 2 minutes. If you want to use the inpaint function, it is best to download and use a machine with more VRAM locally. \
+                    resize_info = gr.Textbox(value="Due to server restrictions, please upload a video that is no longer than 2 minutes. If you want to use the inpaint function, it is best to download 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.")
                     resize_ratio_slider = gr.Slider(minimum=0.02, maximum=1, step=0.02, value=1, label="Resize ratio", visible=True)
           
@@ -534,7 +564,7 @@ with gr.Blocks() as iface:
         # cache_examples=True,
     ) 
 iface.queue(concurrency_count=1)
-iface.launch(debug=True, enable_queue=True)
+iface.launch(debug=True)
 
 
     

inpainter/base_inpainter.py

@@ -7,6 +7,8 @@ import yaml
 import cv2
 import importlib
 import numpy as np
+from tqdm import tqdm
+
 from inpainter.util.tensor_util import resize_frames, resize_masks
 
 
@@ -66,15 +68,15 @@ class BaseInpainter:
         if ratio == 1:
             size = None
         else:
-            size = (int(W*ratio), int(H*ratio))
+            size = [int(W*ratio), int(H*ratio)]
             if size[0] % 2 > 0:
                 size[0] += 1
             if size[1] % 2 > 0:
                 size[1] += 1
         
         masks = np.expand_dims(masks, axis=3)    # expand to T, H, W, 1
-        binary_masks = resize_masks(masks, size)
-        frames = resize_frames(frames, size)          # T, H, W, 3
+        binary_masks = resize_masks(masks, tuple(size))
+        frames = resize_frames(frames, tuple(size))          # T, H, W, 3
         # frames and binary_masks are numpy arrays
 
         h, w = frames.shape[1:3]
@@ -87,7 +89,7 @@ class BaseInpainter:
         imgs, masks = imgs.to(self.device), masks.to(self.device)
         comp_frames = [None] * video_length
 
-        for f in range(0, video_length, self.neighbor_stride):
+        for f in tqdm(range(0, video_length, self.neighbor_stride), desc='Inpainting image'):
             neighbor_ids = [
                 i for i in range(max(0, f - self.neighbor_stride),
                                 min(video_length, f + self.neighbor_stride + 1))

inpainter/model/modules/tfocal_transformer_hq.py

@@ -128,8 +128,10 @@ def window_partition(x, window_size):
         windows: (B*num_windows, T*window_size*window_size, C)
     """
     B, T, H, W, C = x.shape
+
     x = x.view(B, T, H // window_size[0], window_size[0], W // window_size[1],
                window_size[1], C)
+
     windows = x.permute(0, 2, 4, 1, 3, 5, 6).contiguous().view(
         -1, T * window_size[0] * window_size[1], C)
     return windows

requirements.txt

@@ -10,10 +10,7 @@ gradio==3.25.0
 opencv-python
 pycocotools
 matplotlib
-onnxruntime
-onnx
-metaseg==0.6.1
 pyyaml
 av
-mmcv-full
-mmengine
+openmim
+tqdm

track_anything.py

@@ -1,4 +1,6 @@
-import PIL 
+import PIL
+from tqdm import tqdm
+
 from tools.interact_tools import SamControler
 from tracker.base_tracker import BaseTracker
 from inpainter.base_inpainter import BaseInpainter
@@ -10,9 +12,12 @@ import argparse
 class TrackingAnything():
     def __init__(self, sam_checkpoint, xmem_checkpoint, e2fgvi_checkpoint, args):
         self.args = args
-        self.samcontroler = SamControler(sam_checkpoint, args.sam_model_type, args.device)
-        self.xmem = BaseTracker(xmem_checkpoint, device=args.device)
-        self.baseinpainter = BaseInpainter(e2fgvi_checkpoint, args.device) 
+        self.sam_checkpoint = sam_checkpoint
+        self.xmem_checkpoint = xmem_checkpoint
+        self.e2fgvi_checkpoint = e2fgvi_checkpoint
+        self.samcontroler = SamControler(self.sam_checkpoint, args.sam_model_type, args.device)
+        self.xmem = BaseTracker(self.xmem_checkpoint, device=args.device)
+        self.baseinpainter = BaseInpainter(self.e2fgvi_checkpoint, args.device) 
     # def inference_step(self, first_flag: bool, interact_flag: bool, image: np.ndarray, 
     #                    same_image_flag: bool, points:np.ndarray, labels: np.ndarray, logits: np.ndarray=None, multimask=True):
     #     if first_flag:
@@ -39,7 +44,7 @@ class TrackingAnything():
         masks = []
         logits = []
         painted_images = []
-        for i in range(len(images)):
+        for i in tqdm(range(len(images)), desc="Tracking image"):
             if i ==0:           
                 mask, logit, painted_image = self.xmem.track(images[i], template_mask)
                 masks.append(mask)
@@ -51,7 +56,6 @@ class TrackingAnything():
                 masks.append(mask)
                 logits.append(logit)
                 painted_images.append(painted_image)
-                print("tracking image {}".format(i))
         return masks, logits, painted_images
     
         

tracker/base_tracker.py

@@ -126,6 +126,7 @@ class BaseTracker:
     def clear_memory(self):
         self.tracker.clear_memory()
         self.mapper.clear_labels()
+        torch.cuda.empty_cache()
 
 
 ##  how to use: