user137927
user137927

Reputation: 367

Dividing a mask into other masks by its elements

I have a python array mask with arbitrary size which includes (0,1,2,3) numbers. I want to divide it into 4 sub-masks that every one of them includes boolean values which the first one has True in the 0-element positions, second one has True in 1-element position and so on. My input array is too large and I need a fast way to do this in python. Here is an example:

mask = [[0 1 1]
        [2 0 3]
        [3,3,2]
        [1,1,1]]

I want to have:

sub_mask[0] = [[1 0 0]
               [0 1 0]
               [0 0 0]
               [0 0 0]]
sub_mask[1] = [[0 1 1]
               [0 0 0]
               [0 0 0]
               [1 1 1]]
sub_mask[2] = [[0 0 0]
               [1 0 0]
               [0 0 1]
               [0 0 0]]
sub_mask[3] = [[0 0 0]
               [0 0 1]
               [1 1 0]
               [0 0 0]]

Upvotes: 0

Views: 326

Answers (2)

PMende
PMende

Reputation: 5470

You can do this via a simple list comprehension:

sub_mask = [
    mask == val
    for val in np.unique(mask)
]

Upvotes: 2

panktijk
panktijk

Reputation: 1614

You can use numpy:

import numpy as np

mask = [[0, 1, 1],
        [2, 0, 3],
        [3, 3, 2],
        [1, 1, 1]]
mask = np.array(mask)

def get_sub_mask(n):
    # Create an empty sub-mask with all 0's
    sub_mask = np.zeros(mask.shape)
    # Set elements to 1 based on their position in the original mask
    sub_mask[np.where(mask == n)] = 1
    return sub_mask

Upvotes: 2

Related Questions