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

beta -version -li

Browse Source
This commit is contained in:
memoryunreal
2023-04-19 11:34:14 +00:00
parent 579a105166
commit d7f2684303
4 changed files with 266 additions and 168 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
track_anything.py
View File

@@ -15,26 +15,26 @@ class TrackingAnything():
self.xmem = BaseTracker(xmem_checkpoint, device=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:
mask, logit, painted_image = self.samcontroler.first_frame_click(image, points, labels, multimask)
return mask, logit, painted_image
# 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:
# mask, logit, painted_image = self.samcontroler.first_frame_click(image, points, labels, multimask)
# return mask, logit, painted_image
if interact_flag:
mask, logit, painted_image = self.samcontroler.interact_loop(image, same_image_flag, points, labels, logits, multimask)
return mask, logit, painted_image
# if interact_flag:
# mask, logit, painted_image = self.samcontroler.interact_loop(image, same_image_flag, points, labels, logits, multimask)
# return mask, logit, painted_image
mask, logit, painted_image = self.xmem.track(image, logit)
return mask, logit, painted_image
# mask, logit, painted_image = self.xmem.track(image, logit)
# return mask, logit, painted_image
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 interact(self, image: np.ndarray, same_image_flag: bool, points:np.ndarray, labels: np.ndarray, logits: np.ndarray=None, multimask=True):
mask, logit, painted_image = self.samcontroler.interact_loop(image, same_image_flag, points, labels, logits, multimask)
return mask, logit, painted_image
# def interact(self, image: np.ndarray, same_image_flag: bool, points:np.ndarray, labels: np.ndarray, logits: np.ndarray=None, multimask=True):
# mask, logit, painted_image = self.samcontroler.interact_loop(image, same_image_flag, points, labels, logits, multimask)
# return mask, logit, painted_image
def generator(self, images: list, template_mask:np.ndarray):
@@ -53,6 +53,7 @@ class TrackingAnything():
masks.append(mask)
logits.append(logit)
painted_images.append(painted_image)
print("tracking image {}".format(i))
return masks, logits, painted_images