Pad output to fixed length if batch size > 1.

Fix #349

Author: mingxingtan

efficientdet/inference.py

@@ -321,6 +321,10 @@ def det_post_process(params: Dict[Any, Any], cls_outputs: Dict[int, tf.Tensor],
         min_score_thresh=min_score_thresh,
         max_boxes_to_draw=max_boxes_to_draw,
         disable_pyfun=params.get('disable_pyfun'))
+    if params['batch_size'] > 1:
+      # pad to fixed length if batch size > 1.
+      padding_size = max_boxes_to_draw - tf.shape(detections)[0]
+      detections = tf.pad(detections, [[0, padding_size], [0, 0]])
     detections_batch.append(detections)
   return tf.stack(detections_batch, name='detections')
 

efficientdet/model_inspect.py

@@ -104,11 +104,11 @@ def __init__(self,
     self.batch_size = batch_size or None
     self.labels_shape = [batch_size, model_config.num_classes]
 
-    width, height = model_config.image_size
+    height, width = model_config.image_size
     if model_config.data_format == 'channels_first':
-      self.inputs_shape = [batch_size, 3, width, height]
+      self.inputs_shape = [batch_size, 3, height, width]
     else:
-      self.inputs_shape = [batch_size, width, height, 3]
+      self.inputs_shape = [batch_size, height, width, 3]
 
     self.model_config = model_config
 
@@ -161,13 +161,20 @@ def saved_model_inference(self, image_path_pattern, output_dir, **kwargs):
     batch_size = self.batch_size or 1
     all_files = list(tf.io.gfile.glob(image_path_pattern))
     print('all_files=', all_files)
-    num_batches = len(all_files) // batch_size
+    num_batches = (len(all_files) + batch_size - 1) // batch_size
 
     for i in range(num_batches):
       batch_files = all_files[i * batch_size: (i + 1) * batch_size]
-      raw_images = [np.array(Image.open(f)) for f in batch_files]
+      height, width = self.model_config.image_size
+      raw_images = [np.array(Image.open(f).resize((width, height)))
+                    for f in batch_files]
+      size_before_pad = len(raw_images)
+      if size_before_pad < batch_size:
+        padding_size = batch_size - size_before_pad
+        raw_images += [np.zeros_like(raw_images[0])] * padding_size
+      
       detections_bs = driver.serve_images(raw_images)
-      for j in range(len(raw_images)):
+      for j in range(size_before_pad):
         img = driver.visualize(raw_images[j], detections_bs[j], **kwargs)
         img_id = str(i * batch_size + j)
         output_image_path = os.path.join(output_dir, img_id + '.jpg')