objdet_roipool.py

from keras.engine.topology import Layer
class RoiPooling(Layer):
  
  def __init__(self, pool_size, **kwargs):
    self.pool_size = pool_size
    super(RoiPooling, self).__init__(**kwargs)
        
  def build(self, input_shape):
    self.num_channels = input_shape[0][3]
    self.num_rois = input_shape[1][1]

  def compute_output_shape(self, input_shape):
    return None, self.num_rois, self.pool_size, self.pool_size, self.num_channels

  def call(self, thex, mask=None):
    featmap = thex[0]
    roi     = thex[1]

    input_shape = K.shape(featmap) # w,h,c

    output = []
    for i in range(self.num_rois): # for each roi
      # convert the predicted value
      x = K.cast(roi[0, i, 0] * K.cast(input_shape[1], 'float32'), 'int32')
      y = K.cast(roi[0, i, 1] * K.cast(input_shape[2], 'float32'), 'int32')
      w = K.cast(roi[0, i, 2] * K.cast(input_shape[1], 'float32'), 'int32')
      h = K.cast(roi[0, i, 3] * K.cast(input_shape[2], 'float32'), 'int32')
      
      featroi = tf.image.resize_images(featmap[:, x:x+w, y:y+h, :], (self.pool_size, self.pool_size))
      output.append(featroi)
      
    final_output = K.concatenate(output, axis=0)
    final_output = K.reshape(final_output, (-1, self.num_rois, self.pool_size, self.pool_size, self.num_channels))
    

    return final_output

# ROI pooling
featroi = RoiPooling(14)([x,z]) # we ask it to output subimages
roipool = Model(inputs=[x], outputs=[featroi])
roipool.summary()