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

none mask

Browse Source
This commit is contained in:
ShangGaoG
2023-04-14 16:39:40 +08:00
parent 5642d3a1ea
commit 46d55a1404
1 changed files with 14 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
tools/interact_tools.py
View File

@@ -45,17 +45,24 @@ class SamControler():
self.sam_controler.set_image(image)
return
def first_frame_click(self, image: np.ndarray, points:np.ndarray, labels: np.ndarray, multimask=True):
def first_frame_click(self, image: np.ndarray, points:np.ndarray, labels: np.ndarray,logits: np.ndarray=None, multimask=True):
'''
it is used in first frame in video
return: mask, logit, painted image(mask+point)
'''
# self.sam_controler.set_image(image)
self.sam_controler.seg_again(image)
prompts = {
'point_coords': points,
'point_labels': labels,
}
self.sam_controler.set_image(image)
if logits is None:
prompts = {
'point_coords': points,
'point_labels': labels,
}
else:
prompts = {
'point_coords': points,
'point_labels': labels,
'mask_input': logits[None, :, :]
}
masks, scores, logits = self.sam_controler.predict(prompts, 'point', multimask)
mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :]