from 640 hard codes

official/vision/beta/projects/yolo/dataloaders/yolo_input.py

@@ -244,7 +244,7 @@ class Parser(parser.Parser):
     else:
       image = tf.image.resize(
           image, (self._image_h, self._image_w), method='nearest')
-      output_size = tf.cast([640, 640], tf.float32)
+      output_size = tf.cast([self._image_h, self._image_w], tf.float32)
       boxes_ = bbox_ops.denormalize_boxes(boxes, output_size)
       inds = bbox_ops.get_non_empty_box_indices(boxes_)
       boxes = tf.gather(boxes, inds)
@@ -286,7 +286,8 @@ class Parser(parser.Parser):
     # Clip and clean boxes.
     image = image / 255.0
     boxes, inds = preprocessing_ops.transform_and_clip_boxes(
-        boxes, infos, shuffle_boxes=False, area_thresh=0.0, augment=True)
+        boxes, infos, shuffle_boxes=False, area_thresh=0.0, augment=True,
+        output_size = [self._image_h, self._image_w])
     classes = tf.gather(classes, inds)
     info = infos[-1]
 
@@ -342,6 +343,9 @@ class Parser(parser.Parser):
 
     # Update the labels dictionary.
     if not is_training:
+      output_size = tf.cast([height, width], tf.float32)
+      boxes = bbox_ops.denormalize_boxes(gt_boxes, output_size)
+      gt_area = (boxes[2] - boxes[0]) * (boxes[3] - boxes[1])
 
       # Sets up groundtruth data for evaluation.
       groundtruths = {
@@ -352,7 +356,7 @@ class Parser(parser.Parser):
           'image_info': info,
           'boxes': gt_boxes,
           'classes': gt_classes,
-          'areas': tf.gather(data['groundtruth_area'], inds),
+          'areas': gt_area,
           'is_crowds':
               tf.cast(tf.gather(data['groundtruth_is_crowd'], inds), tf.int32),
       }

official/vision/beta/projects/yolo/dataloaders/yolo_input_test.py

+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Yolo Dataset Testing functions"""
+from official.vision.beta.projects.yolo.common import registry_imports  # pylint: disable=unused-import
+from official.vision.beta.projects.yolo.tasks import image_classification as imc
+from official.vision.beta.projects.yolo.configs import darknet_classification as dcfg
+
+import os
+import tensorflow as tf
+from official.core import train_utils
+from official.core import task_factory
+from absl.testing import parameterized
+
+PATH_TO_COCO = '/media/vbanna/DATA_SHARE/CV/datasets/COCO_raw/records/'
+
+def test_yolo_input_task(scaled_pipeline = True, batch_size = 1):
+  if not scaled_pipeline:
+    experiment = "yolo_darknet"
+    config_path = [
+      "official/vision/beta/projects/yolo/configs/experiments/yolov4/tpu/512.yaml"]
+  else:
+    experiment = "large_yolo"
+    config_path = [
+      "official/vision/beta/projects/yolo/configs/experiments/scaled-yolo/detection/yolo_l_p6_1280_tpu.yaml"]
+
+  config = train_utils.ParseConfigOptions(experiment=experiment, 
+                                          config_file=config_path)
+  params = train_utils.parse_configuration(config)
+  config = params.task
+  task = task_factory.get_task(params.task)
+
+  config.train_data.global_batch_size = batch_size
+  config.validation_data.global_batch_size = 1
+  config.train_data.dtype = 'float32'
+  config.validation_data.dtype = 'float32'
+  config.validation_data.shuffle_buffer_size = 1
+  config.train_data.shuffle_buffer_size = 1
+  config.train_data.input_path = os.path.join(PATH_TO_COCO, 'train*')
+  config.validation_data.input_path = os.path.join(PATH_TO_COCO, 'val*')
+
+  with tf.device('/CPU:0'):
+    train_data = task.build_inputs(config.train_data)
+    test_data = task.build_inputs(config.validation_data)
+  return train_data, test_data, config
+
+def test_yolo_pipeline_visually(is_training=True, num=30):
+  # visualize the datapipeline
+  import matplotlib.pyplot as plt
+  dataset, testing, _ = test_yolo_input_task()
+
+  data = dataset if is_training else testing
+  data = data.take(num)
+  for l, (image, label) in enumerate(data):
+    image = tf.image.draw_bounding_boxes(
+        image, label['bbox'], [[1.0, 1.0, 1.0]])
+
+    gt = label['true_conf']
+
+    obj3 = tf.clip_by_value(gt['3'][..., 0], 0.0, 1.0)
+    obj4 = tf.clip_by_value(gt['4'][..., 0], 0.0, 1.0)
+    obj5 = tf.clip_by_value(gt['5'][..., 0], 0.0, 1.0)
+    obj6 = tf.clip_by_value(gt['6'][..., 0], 0.0, 1.0)
+
+    for shind in range(1):
+      fig, axe = plt.subplots(1, 5)
+
+      image = image[shind]
+
+      axe[0].imshow(image)
+      axe[1].imshow(obj3[shind, ..., :3].numpy())
+      axe[2].imshow(obj4[shind, ..., :3].numpy())
+      axe[3].imshow(obj5[shind, ..., :3].numpy())
+      axe[4].imshow(obj6[shind, ..., :3].numpy())
+
+      fig.set_size_inches(18.5, 6.5, forward=True)
+      plt.tight_layout()
+      plt.show()
+
+class YoloDetectionInputTest(tf.test.TestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(('scaled', True), ('darknet', False))
+  def test_yolo_input(self, scaled_pipeline):
+    # builds a pipline forom the config and tests the datapipline shapes
+    dataset, _, params = test_yolo_input_task(
+        scaled_pipeline=scaled_pipeline, 
+        batch_size=1)
+
+    dataset = dataset.take(1)
+
+    for image, label in dataset:
+      self.assertAllEqual(image.shape, ([1] + params.model.input_size))
+      self.assertTrue(
+          tf.reduce_all(tf.math.logical_and(image >= 0, image <= 1)))
+
+
+if __name__ == '__main__':
+  # tf.test.main()
+  test_yolo_pipeline_visually(is_training=True, num=20)
\ No newline at end of file

official/vision/beta/projects/yolo/ops/mosaic.py

@@ -180,7 +180,8 @@ class Mosaic:
         area_thresh=self._area_thresh,
         shuffle_boxes=False,
         augment=True,
-        seed=self._seed)
+        seed=self._seed,
+        output_size=[self._output_size[0], self._output_size[1]])
     classes, is_crowd, area = self._select_ind(inds, classes, is_crowd, area)  # pylint:disable=unbalanced-tuple-unpacking
     return image, boxes, classes, is_crowd, area, crop_points
 
@@ -229,7 +230,8 @@ class Mosaic:
         None,
         affine=affine,
         area_thresh=self._area_thresh,
-        seed=self._seed)
+        seed=self._seed,
+        output_size=[self._output_size[0], self._output_size[1]])
     classes, is_crowd, area = self._select_ind(inds, classes, is_crowd, area)  # pylint:disable=unbalanced-tuple-unpacking
     return image, boxes, classes, is_crowd, area, area
 

official/vision/beta/projects/yolo/ops/preprocessing_ops.py

@@ -839,7 +839,8 @@ def transform_and_clip_boxes(boxes,
                              shuffle_boxes=False,
                              area_thresh=0.1,
                              seed=None,
-                             augment=True):
+                             augment=True,
+                             output_size = None):
   """Clips and cleans the boxes.
 
   Args:
@@ -870,7 +871,10 @@ def transform_and_clip_boxes(boxes,
 
   # Make sure all boxes are valid to start, clip to [0, 1] and get only the
   # valid boxes.
-  output_size = tf.cast([640, 640], tf.float32)
+  if output_size is None:
+    output_size = tf.cast([640, 640], tf.float32)
+  else:
+    output_size = tf.cast(output_size, tf.float32)
   if augment:
     boxes = tf.math.maximum(tf.math.minimum(boxes, 1.0), 0.0)
   cond = get_valid_boxes(boxes)