title: Vectorising the python code for modifying masks used in Image segmentation
body: <p>I am working on an Image segmentation challenge. I have masks which have 5 labels(0, 1, 2, 3, 4) and the layout of one such mask(2D matrix) is : </p>
  
  <pre><code>[0 0 0 0 0 0
   0 0 1 1 2 2
   1 1 1 1 2 2
   2 2 2 2 2 2
   3 3 3 3 3 3
   4 4 4 4 4 4]
  </code></pre>
  
  <p>I want to merge several classes such that the modified mask looks like :</p>
  
  <pre><code>[0 0 0 0 0 0
   0 0 0 0 1 1
   0 0 0 0 1 1
   1 1 1 1 1 1
   2 2 2 2 2 2
   2 2 2 2 2 2]
  </code></pre>
  
  <p>Merged 0 and 1 to 0.
  Changed 2 to 1.
  Merged 3 and 4, changed them to 2.</p>
  
  <p>I have implemented a looped version and it is taking a lot of time as the dimension of my mask is (601, 462, 951).</p>
  
  <pre><code>for i in range(0, dim.shape[0]):
    for j in range(0, dim.shape[1]):
      for k in range(0, dim.shape[2]):
        if dim[i, j, k] in (0, 2):
          dim[i, j, k] = 1
  
        if dim[i, j, k] == 3:
          dim[i, j, k] = 2
  
        if dim[i, j, k] in (4, 8):
          dim[i, j, k] = 3
  
        if dim[i, j, k] == 9:
          dim[i, j, k] = 4
  </code></pre>
  
  <p>I can't find any way to vectorize my code so that I can remove the loops.</p>