Python numpy array assignment to integer indexed flat slice
<p>While learning numpy, I wrote code which doing LSB(steganography) encryption:</p>

<pre><code>def str2bits_nparray(s):
    return np.array(map(int, (''.join(map('{:07b}'.format, bytearray(s))))), dtype=np.bool)

def LSB_encode(img, msg, channel):
    msg_bits = str2bits_nparray(msg)
    xor_mask = np.zeros_like(img, dtype=np.bool)
    xor_mask[:, :, channel].flat[:len(msg_bits)] = np.ones_like(msg_bits, dtype=np.bool)
    img[xor_mask] = img[xor_mask] &gt;&gt; 1 &lt;&lt; 1 | msg_bits


msg = 'A' * 1000
img_name = 'screenshot.png'
chnl = 2
img = imread(img_name)
LSB_encode(img, msg, chnl)
</code></pre>

<p>Code works fine, but when i'm trying to made <code>chnl = [2, 1]</code> this line:</p>

<pre><code>xor_mask[:, :, channel].flat[:len(msg_bits)] = np.ones_like(msg_bits, dtype=np.bool)
</code></pre>

<p>doesnt assign value to <code>xor_mask</code> with </p>

<p><code>xor_mask[:, :,</code><strong>[2, 1]</strong><code>].flat[:len(msg_bits)]</code></p>

<p>Is there way to fix this?</p>

<p>I tryed solution with for-loop over channels:</p>

<pre><code>for ch in channel:
    xor_mask[:, :, ch].flat[:len(msg_bits)] = np.ones_like(msg_bits, dtype=np.bool)
</code></pre>

<p>But this is doing not that i want from </p>

<p><code>xor_mask[:, :,</code><strong>[2, 1]</strong><code>].flat[:len(msg_bits)] = np.ones_like(msg_bits, dtype=np.bool)</code></p>

<p>IIUC, here's an approach to get the linear indices, then slice to the length of no. of elems required to be set and then perform the setting -</p>

<pre><code>m,n,r = xor_mask.shape  # Store shape info

# Create range arrays corresponding to those shapes
x,y,z = np.ix_(np.arange(m),np.arange(n),channel)

# Get the indices to be set and finaally perform the setting
idx = (x*n*r + y*r + z).ravel()[:len(msg_bits)]
xor_mask.ravel()[idx] = 1
</code></pre>

<p>Sample run -</p>

<pre><code>In [180]: xor_mask
Out[180]: 
array([[[25, 84, 37, 96, 72, 84, 91],
        [94, 56, 78, 71, 48, 65, 98]],

       [[33, 56, 14, 92, 90, 64, 76],
        [71, 71, 77, 31, 96, 36, 49]]])

In [181]: # Other inputs
     ...: channel = np.array([2,1])
     ...: msg_bits = np.array([2,3,6,1,4])
     ...: 

In [182]: m,n,r = xor_mask.shape  # Store shape info
     ...: x,y,z = np.ix_(np.arange(m),np.arange(n),channel)
     ...: idx = (x*n*r + y*r + z).ravel()[:len(msg_bits)]
     ...: xor_mask.ravel()[idx] = 1
     ...: 

In [183]: xor_mask # First 5 elems from flattend version
                   # of xor_mask[:,:,channel] set as 1 
                   # as len(msg_bits) = 5.
Out[183]: 
array([[[25,  1,  1, 96, 72, 84, 91],
        [94,  1,  1, 71, 48, 65, 98]],

       [[33, 56,  1, 92, 90, 64, 76],
        [71, 71, 77, 31, 96, 36, 49]]])
</code></pre>

<p>Instead, if you were trying to set for all elems across all dimensions in <code>3D</code> input array along the first of <code>channel</code> : <code>2</code> and then along the second one <code>1</code> and so on, we need to create <code>idx</code> differently, like so -</p>

<pre><code>idx = (x*n*r + y*r + z).transpose(2,0,1).ravel()[:len(msg_bits)]
</code></pre>

python arrays python-2.7 numpy slice