Mirrors/Track-Anything
1
0
Fork 0
mirror of https://github.com/gaomingqi/Track-Anything.git synced 2026-05-18 13:15:51 +02:00
Code Issues Packages Projects Releases Wiki Activity

gaoshang

Browse Source
This commit is contained in:
ShangGaoG
2023-04-17 22:59:49 +08:00
parent 3053948124 5ee4c5cadb
commit 20b81c49c3
3 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tools/base_segmenter.py
View File

@@ -85,7 +85,7 @@ if __name__ == "__main__":
# initialise BaseSegmenter
SAM_checkpoint= '/ssd1/gaomingqi/checkpoints/sam_vit_h_4b8939.pth'
model_type = 'vit_h'
device = "cuda:0"
device = "cuda:4"
base_segmenter = BaseSegmenter(SAM_checkpoint=SAM_checkpoint, model_type=model_type, device=device)
# image embedding (once embedded, multiple prompts can be applied)

2
track_anything.py
View File

@@ -12,7 +12,7 @@ class TrackingAnything():
def __init__(self, sam_checkpoint, xmem_checkpoint, args):
self.args = args
self.samcontroler = SamControler(sam_checkpoint, args.sam_model_type, args.device)
self.xmem = BaseTracker(xmem_checkpoint, device=args.device, sam_model=sam_checkpoint, model_type=args.sam_model_type)
self.xmem = BaseTracker(xmem_checkpoint, device=args.device)
def inference_step(self, first_flag: bool, interact_flag: bool, image: np.ndarray,

2
tracker/base_tracker.py
View File

@@ -20,7 +20,7 @@ from torchvision.transforms import Resize
class BaseTracker:
def __init__(self, xmem_checkpoint, device, sam_model, model_type=None) -> None:
def __init__(self, xmem_checkpoint, device, sam_model=None, model_type=None) -> None:
"""
device: model device
xmem_checkpoint: checkpoint of XMem model