Mirrors/Track-Anything
1
0
Fork 0
mirror of https://github.com/gaomingqi/Track-Anything.git synced 2025-12-16 16:37:58 +01:00
Code Issues Packages Projects Releases Wiki Activity

update base_tracker

Browse Source
This commit is contained in:
gaomingqi
2023-04-14 22:17:47 +08:00
parent 2809c45402
commit 746d276121
1 changed files with 5 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tracker/base_tracker.py
View File

@@ -1,10 +1,9 @@
# input: frame list, first frame mask # import for debugging
# output: segmentation results on all frames
import os import os
import glob import glob
import numpy as np import numpy as np
from PIL import Image from PIL import Image
# import for base_tracker
import torch import torch
import yaml import yaml
import torch.nn.functional as F import torch.nn.functional as F
@@ -13,8 +12,6 @@ from inference.inference_core import InferenceCore
from util.mask_mapper import MaskMapper from util.mask_mapper import MaskMapper
from torchvision import transforms from torchvision import transforms
from util.range_transform import im_normalization from util.range_transform import im_normalization
import sys import sys
sys.path.insert(0, sys.path[0]+"/../") sys.path.insert(0, sys.path[0]+"/../")
from tools.painter import mask_painter from tools.painter import mask_painter
@@ -40,6 +37,7 @@ class BaseTracker:
]) ])
self.device = device self.device = device
# changable properties
self.mapper = MaskMapper() self.mapper = MaskMapper()
self.initialised = False self.initialised = False
@@ -109,7 +107,7 @@ if __name__ == '__main__':
# ---------------------------------------------------------- # ----------------------------------------------------------
device = 'cuda:1' device = 'cuda:1'
XMEM_checkpoint = '/ssd1/gaomingqi/checkpoints/XMem-s012.pth' XMEM_checkpoint = '/ssd1/gaomingqi/checkpoints/XMem-s012.pth'
tracker = BaseTracker(device, XMEM_checkpoint) tracker = BaseTracker(XMEM_checkpoint, device)
# track anything given in the first frame annotation # track anything given in the first frame annotation
for ti, frame in enumerate(frames): for ti, frame in enumerate(frames):