Not working while passing multiple prompt points

Nobody has claimed this yet.

For specifying a specific object with multiple points, I am seeing an error while passing more than two points.

#Cell 61
input_point = np.array([[500, 1375], [1500, 1375], [1125, 1625]])
input_label = np.array([1, 1])

mask_input = logits[np.argmax(scores), :, :] # Choose the model's best mask

#Cell 62
masks, _, _ = predictor.predict(
point_coords=input_point,
point_labels=input_label,
mask_input=mask_input[None, :, :],
multimask_output=False,
)

--------------------------------------------------------------------------- IndexError Traceback (most recent call last) Cell In[68], line 1
----> 1 masks, _, _ = predictor.predict(
2 point_coords=input_point,
3 point_labels=input_label,
4 mask_input=mask_input[None, :, :],
5 multimask_output=False,
6 )

/segment-anything/segment_anything/predictor.py:154), in SamPredictor.predict(self, point_coords, point_labels, box, mask_input, multimask_output, return_logits)
151 mask_input_torch = torch.as_tensor(mask_input, dtype=torch.float, device=self.device)
152 mask_input_torch = mask_input_torch[None, :, :, :]
--> 154 masks, iou_predictions, low_res_masks = self.predict_torch(
155 coords_torch,
156 labels_torch,
157 box_torch,
158 mask_input_torch,
159 multimask_output,
160 return_logits=return_logits,
161 )
163 masks_np = masks[0].detach().cpu().numpy()
164 iou_predictions_np = iou_predictions[0].detach().cpu().numpy()

/python3.9/site-packages/torch/autograd/grad_mode.py:28), in _DecoratorContextManager.__call__..decorate_context(*args, **kwargs)
25 @functools.wraps(func)
26 def decorate_context(*args, **kwargs):
27 with self.__class__():
---> 28 return func(*args, **kwargs)

/segment-anything/segment_anything/predictor.py:222), in SamPredictor.predict_torch(self, point_coords, point_labels, boxes, mask_input, multimask_output, return_logits)
219 points = None
221 # Embed prompts
--> 222 sparse_embeddings, dense_embeddings = self.model.prompt_encoder(
223 points=points,
224 boxes=boxes,
225 masks=mask_input,
226 )
228 # Predict masks
229 low_res_masks, iou_predictions = self.model.mask_decoder(
230 image_embeddings=self.features,
231 image_pe=self.model.prompt_encoder.get_dense_pe(), (...)
234 multimask_output=multimask_output,
235 )

python3.9/site-packages/torch/nn/modules/module.py:1102), in Module._call_impl(self, *input, **kwargs)
1098 # If we don't have any hooks, we want to skip the rest of the logic in
1099 # this function, and just call forward.
1100 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks 1101 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1102 return forward_call(*input, **kwargs)
1103 # Do not call functions when jit is used
1104 full_backward_hooks, non_full_backward_hooks = [], []

projects/segment-anything/segment_anything/modeling/prompt_encoder.py:155), in PromptEncoder.forward(self, points, boxes, masks)
153 if points is not None:
154 coords, labels = points
--> 155 point_embeddings = self._embed_points(coords, labels, pad=(boxes is None))
156 sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
157 if boxes is not None:

segment-anything/segment_anything/modeling/prompt_encoder.py:87), in PromptEncoder._embed_points(self, points, labels, pad)
85 labels = torch.cat([labels, padding_label], dim=1)
86 point_embedding = self.pe_layer.forward_with_coords(points, self.input_image_size)
---> 87 point_embedding[labels == -1] = 0.0
88 point_embedding[labels == -1] += self.not_a_point_embed.weight
89 point_embedding[labels == 0] += self.point_embeddings[0].weight

IndexError: The shape of the mask [1, 3] at index 1 does not match the shape of the indexed tensor [1, 4, 256] at index 1

Could you please check this bug?

Reproduce the failure from the inputs in notebook cells 61–62, then read segment_anything/predictor.py and segment_anything/modeling/prompt_encoder.py at the traceback lines. Compare the point and label shapes and inspect existing examples or tests for prompt handling. Done means multiple point prompts run without an indexing error and the behavior has regression coverage.

Written by the indexing model from the issue text.