title: How can I get a list of pixels of specific class from masked image using tensorflow on GPU?
body: <p>I have a segmented image (Ex: 200,200,20 -&gt; 20 classes) and an image(200,200,3). I'm trying to get the color of list of pixels belonging to a particular class.</p>
  <p>I'm able to accomplish this task using numpy operations but this pixel level computation is very costly. So I'm trying to create a tensorflow model to run it on GPU, that takes input image and segmented image. Then trying to get the list of pixels which I will input to a function that i have which returns the most dominant color. But first i need to get the list of pixels.</p>
  <p>I did it with numpy and almost with tensorflow. But with tensorflow I'm not able to accomplish getting the list of pixels for batch of images. For 1 image its working but for batch of images, at the stage <code>pixel_list = tf.boolean_mask(orig_img,mask)</code> the pixel list is being combined together so it's not possible to differentiate pixels that belong to which image.</p>
  <p><strong>Example</strong>: 4 RGB images, 20 classes</p>
  <p>inputs to model:
  image_tensor = 4,200,200,3
  , segmented_image = 4,200,200,20</p>
  <p>At the stage <code>pixel_list = tf.boolean_mask(orig_img,mask)</code></p>
  <p>orig_image: (4,200,200,3)</p>
  <p>mask: (4, 200,200)</p>
  <p>expectation: (4,) -&gt; [ (n1,3), (n2,3), (n3,3), (n4,3) ] {variable, ni is number of pixels belonging to a particular class in image i}</p>
  <p>Instead I'm getting: <code>(n1+n2+n3+n4),3</code>, the dimensions are lost, pixel list is being combined together. So its not possible to differentiate pixels that belong to which image.</p>
  <p>I'm also open to other ways of getting color of segmented classes if you can suggest.</p>
  <p>code for reference:</p>
  <pre class="lang-py prettyprint-override"><code>def create_model():
      segmented_image = tf.keras.Input(shape=(200, 200, 1),
              dtype=tf.float32, name='post_input')
      original_image = tf.keras.Input(shape=(200, 200, 3),
                                      dtype=tf.float32, name='image')
      recognized_classes = [
          'hat',
          'upper_clothes',
          'dress',
          'coat',
          'pants',
          'jumpsuits',
          'skirt',
          ]
  
      # currently supported classes
  
      classes = [
          1,
          5,
          6,
          7,
          9,
          10,
          12,
          ]
  
      # mask and pixel list for class 1
  
      pixel_list_result = dict()
      placer = 0
      for i in classes:
          print '=-' * 50
          print ('for ', recognized_classes[placer])
  
          # # Creating mask of class i
  
          mask = tf.cast(tf.math.equal(parse_result, i), tf.bool)  # none,240,240
  
          # # Logits
  
          logit_list = tf.boolean_mask(logit_result, mask)
          logit_list = tf.math.sigmoid(logit_list)
          logit_mean = tf.math.reduce_mean(logit_list, axis=0,
                  name=recognized_classes[placer])
  
          # # Pixels list
  
          pixel_list = tf.boolean_mask(orig_img, mask)
          pixel_list = tf.reverse(pixel_list, axis=[-1],
                                  name=recognized_classes[placer])
  
          # # Final
  
          pixel_list_result[recognized_classes[placer]] = [pixel_list,
                  logit_mean]
          placer += 1
  
      return tf.keras.Model(inputs=[model_output, im_input], outputs=[
          pixel_list_result['hat'],
          pixel_list_result['upper_clothes'],
          pixel_list_result['dress'],
          pixel_list_result['coat'],
          pixel_list_result['pants'],
          pixel_list_result['jumpsuits'],
          pixel_list_result['skirt'],
          ], name='postprocessing')
  </code></pre>