Solved: Grabbing Individual Masks from `get_preds`

fastai
1

Hi all,

I’m trying to grab individual masks from get_preds.

I know I can achieve this with one mask by doing:

preds = learn.get_preds()
t = preds[0].argmax(dim=1)
learn.data.test_ds.y.reconstruct(t)

But in bulk predictions I’m not finding a good way to split those masks. Any thoughts?

Eg:

preds = learn.get_preds()
t = preds[0][0].argmax(dim=1)
learn.data.test_ds.y.reconstruct(t)

will not work. Thoughts?

Thanks!

2

Solved:

Those tensors need to be unsqueezed first.

t = preds[0][0].unsqueeze(dim=0).argmax(dim=1)

1 Like
3

Well alternatively you can do t = preds[0][0].argmax(dim=0) which has shape [H, W]. If you want something of shape [1, H, W], you can do t = preds[0][0].argmax(dim=0, keepdim=True). Your solution is absolutely fine though, just wasn’t my first intuition.

EDIT: looking back, your solution is better as it makes more sense to always have the same prediction pipeline, which here is using argmax(dim=1).

1 Like