batched inference and feature extraction for point prompts

[ysig]

Hi,

Thank you for this great work. I have a question:
Given a batched input x: B,3,H,W, what is the best way to:

a) batch preprocess the images (so that you can offload this to a dataloader)
b) do a batch segmentation given a prompt: B, K, 2 and label: B, K, 2
c) extract both the mask and the features of the mask (can be an average) for the resulting segmentation?

Thank you,