Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,885 Bytes
59b2a81 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import os, sys
import PIL
from tqdm import tqdm
import numpy as np
import argparse
# Import files from the local folder
root_path = os.path.abspath('.')
sys.path.append(root_path)
from track_anything_code.tools.interact_tools import SamControler
from track_anything_code.tracker.base_tracker import BaseTracker
class TrackingAnything():
def __init__(self, sam_checkpoint, xmem_checkpoint, args):
self.args = args
self.sam_checkpoint = sam_checkpoint
self.xmem_checkpoint = xmem_checkpoint
self.samcontroler = SamControler(self.sam_checkpoint, args.sam_model_type, args.device)
self.xmem = BaseTracker(self.xmem_checkpoint, device=args.device)
def first_frame_click(self, image: np.ndarray, points:np.ndarray, labels: np.ndarray, multimask=True):
mask, logit, painted_image = self.samcontroler.first_frame_click(image, points, labels, multimask)
return mask, logit, painted_image
def generator(self, images: list, template_mask:np.ndarray):
masks = []
logits = []
painted_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)
logits.append(logit)
painted_images.append(painted_image)
else:
mask, logit, painted_image = self.xmem.track(images[i])
masks.append(mask)
logits.append(logit)
painted_images.append(painted_image)
return masks, logits, painted_images
def parse_augment():
parser = argparse.ArgumentParser()
parser.add_argument('--device', type=str, default="cuda:0")
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=False)
args = parser.parse_args()
if args.debug:
print(args)
return args
# if __name__ == "__main__":
# masks = None
# logits = None
# painted_images = None
# images = []
# image = np.array(PIL.Image.open('/hhd3/gaoshang/truck.jpg'))
# args = parse_augment()
# # images.append(np.ones((20,20,3)).astype('uint8'))
# # images.append(np.ones((20,20,3)).astype('uint8'))
# images.append(image)
# images.append(image)
# mask = np.zeros_like(image)[:,:,0]
# mask[0,0]= 1
# trackany = TrackingAnything('/ssd1/gaomingqi/checkpoints/sam_vit_h_4b8939.pth','/ssd1/gaomingqi/checkpoints/XMem-s012.pth', args)
# masks, logits ,painted_images= trackany.generator(images, mask)
|