Batched predictor inference returns inconsistently shaped masks

[geodavic]

When performing batch inference using SAM2ImagePredictor, if you pass a singleton batch of points/labels, you get masks of shape (3, H,W):

>>> predictor.set_image(image)
>>> point_coords.shape # Singleton batch
(1, 2, 2)
>>> masks, _ , _ = predictor.predict(point_coords=point_coords, point_label=point_labels)
>>> masks.shape
(3, 1200,1800)

However, when your batch size is N>1 your masks become shape (N, 3, H, W). This makes it difficult to use the predictor interface where the batch size is arbitrary (and user-specified).