Faster segmentation processing in dataloader

Hi,

I was playing around with the code a little and noticed that the processing of the segmentation mask in the dataloader is fairly slow. (Line 73 - 82 in dataset.py). I quickly threw together a version of this which is much faster. I propose to replace

count = 0
for i in range(mask.shape[0]):
    for j in range(mask.shape[1]):
        if mask[i,j] not in mapping:
            if mask[mask == mask[i,j]].sum()>50:
                mapping[mask[i,j]] = count
                count += 1
            else:
                mapping[mask[i,j]] = 0
        mask[i,j] = mapping[mask[i,j]]

with

values, indices, counts = np.unique(mask, return_inverse=True, return_counts=True)
to_eliminate = counts <= 50
mapping = np.arange(len(values))
mapping[to_eliminate] = 0
_h, _w = mask.shape
mask = mapping[indices].reshape((_h, _w))

On my machine, this sections gets speed up from ~0.3s to 0.01s.

Best, Felix

zpbao / Discovery_Obj_Move

Faster segmentation processing in dataloader #4