[Python] The issue of mask fragmentation during SAM2 tracking

I am currently working on object tracking. I use Moondream2 to identify objects in the scene, filter out duplicate bounding boxes, and then use SAM2 to track the objects. During the tracking process, the initial mask might be incomplete, which can lead to fragmented masks. To improve tracking accuracy, I want to implement a method in my code to determine whether a mask is continuous or fragmented. How can I achieve this?

start_flag = time.time()
# Visual Querying
scan_text_prompt2 = 'Can you describe the cash register screen and everything on the counter? Please list all the items on the counter, excluding the screen.'
scan_ans2:str = model.query(crop_PIL, f"{scan_text_prompt2}")["answer"]
print(f"\nVisual query: '{scan_text_prompt2}'")
print(scan_ans2)
title_text += f'{scan_text_prompt2} {scan_ans2}\n'

end_flag = time.time()
print('Process time: ' + str(round(end_flag - start_flag, 3)) + 'seconds')


title_text = ''
start_flag = time.time()


locate_text_prompt = scan_ans2
print(f"\nObject detection: {locate_text_prompt}")


detect_prompts = [prompt.strip() for prompt in locate_text_prompt.split(',')]

objects = []

for prompt in detect_prompts:
print(f"Detecting: {prompt}")
detection_result = model.detect(crop_PIL, prompt)["objects"]
print(f"Found {len(detection_result)} product(s) for '{prompt}'")

objects.extend(detection_result)


title_text += f'{prompt}:\n'


print("Final detected objects:", objects)

end_flag = time.time()
print('Process time: ' + str(round(end_flag - start_flag, 3)) + ' seconds')

filtered_boxes = remove_overlapping_boxes(objects, iou_threshold=0.3)

print(filtered_boxes)

for i, box in enumerate(filtered_boxes):
print(f"Box {i}: {box}")
image_np = cv2.cvtColor(np.array(crop_PIL.copy()), cv2.COLOR_BGR2RGB)


for bbox in filtered_boxes:
result = draw_bbox(image_np, bbox)

try:
for bbox in filtered_boxes:
result = draw_bbox(image_np, bbox)
plt.imshow(cv2.cvtColor(np.array(result), cv2.COLOR_BGR2RGB))
#plt.title(title_text)
except:
plt.imshow(cv2.cvtColor(np.array(image_np), cv2.COLOR_BGR2RGB))
#plt.title(title_text)
print("error")

filtered_boxes_np = []
for bbox in filtered_boxes:
result = box_dict_to_array(image_np, filtered_boxes)
filtered_boxes_np.append(result)
print(filtered_boxes_np)




ann_frame_idx = 0


prompts = {}


plt.figure(figsize=(9, 6))
plt.title(f"frame {ann_frame_idx}")
plt.imshow(Image.open(os.path.join(video_dir, frame_names[ann_frame_idx])))


for i, box in enumerate(filtered_boxes_np):
ann_obj_id = i



points = get_two_pt(box)
labels = np.array([1, 1], np.int32)


prompts[ann_obj_id] = points, labels


_, out_obj_ids, out_mask_logits = predictor.add_new_points_or_box(
inference_state=inference_state,
frame_idx=ann_frame_idx,
obj_id=ann_obj_id,
points=points,
labels=labels,
)


show_points(points, labels, plt.gca())
for j, out_obj_id in enumerate(out_obj_ids):
show_points(*prompts[out_obj_id], plt.gca())
show_mask((out_mask_logits[j] > 0.0).cpu().numpy(), plt.gca(), obj_id=out_obj_id)

plt.axis('off')
plt.tight_layout()
plt.show()


after doing this

print(out_mask_logits[0])
plt.imshow(out_mask_logits[0][0].cpu().numpy(), cmap='viridis')
plt.colorbar()
plt.title("Logit Heatmap")
plt.show()


I got

tensor([[[-13.0000, -13.0000, -13.0000, ..., -16.2500, -16.2500, -16.2500],
[-13.1292, -13.1292, -13.1292, ..., -16.0312, -16.0312, -16.0312],
[-14.5069, -14.5069, -14.5069, ..., -13.6979, -13.6979, -13.6979],
...,
[-12.0451, -12.0451, -12.0451, ..., -11.9167, -11.9167, -11.9167],
[-11.4896, -11.4896, -11.4896, ..., -11.2500, -11.2500, -11.2500],
[-11.4375, -11.4375, -11.4375, ..., -11.1875, -11.1875, -11.1875]]],
device='cuda:0')




How do I solve the situation



I would like to have a condition or method to help me distinguish broken (disconnected) mask regions, instead of using a for loop to check them one by one.

Continue reading...