Merge pull request #10263 from PurdueDualityLab:dataload_pr

PiperOrigin-RevId: 399483092

official/vision/beta/projects/yolo/ops/preprocess_ops_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.
-
-"""preprocess_ops tests."""
-from absl.testing import parameterized
-import numpy as np
-import tensorflow as tf
-
-from official.vision.beta.projects.yolo.ops import preprocess_ops
-
-
-class PreprocessOpsTest(parameterized.TestCase, tf.test.TestCase):
-
-  @parameterized.parameters((416, 416, 5, 300, 300), (100, 200, 6, 50, 50))
-  def test_resize_crop_filter(self, default_width, default_height, num_boxes,
-                              target_width, target_height):
-    image = tf.convert_to_tensor(
-        np.random.rand(default_width, default_height, 3))
-    boxes = tf.convert_to_tensor(np.random.rand(num_boxes, 4))
-    resized_image, resized_boxes = preprocess_ops.resize_crop_filter(
-        image, boxes, default_width, default_height, target_width,
-        target_height)
-    resized_image_shape = tf.shape(resized_image)
-    resized_boxes_shape = tf.shape(resized_boxes)
-    self.assertAllEqual([default_height, default_width, 3],
-                        resized_image_shape.numpy())
-    self.assertAllEqual([num_boxes, 4], resized_boxes_shape.numpy())
-
-  @parameterized.parameters((7, 7., 5.), (25, 35., 45.))
-  def test_translate_boxes(self, num_boxes, translate_x, translate_y):
-    boxes = tf.convert_to_tensor(np.random.rand(num_boxes, 4))
-    translated_boxes = preprocess_ops.translate_boxes(
-        boxes, translate_x, translate_y)
-    translated_boxes_shape = tf.shape(translated_boxes)
-    self.assertAllEqual([num_boxes, 4], translated_boxes_shape.numpy())
-
-  @parameterized.parameters((100, 200, 75., 25.), (400, 600, 25., 75.))
-  def test_translate_image(self, image_height, image_width, translate_x,
-                           translate_y):
-    image = tf.convert_to_tensor(np.random.rand(image_height, image_width, 4))
-    translated_image = preprocess_ops.translate_image(
-        image, translate_x, translate_y)
-    translated_image_shape = tf.shape(translated_image)
-    self.assertAllEqual([image_height, image_width, 4],
-                        translated_image_shape.numpy())
-
-  @parameterized.parameters(([1, 2], 20, 0), ([13, 2, 4], 15, 0))
-  def test_pad_max_instances(self, input_shape, instances, pad_axis):
-    expected_output_shape = input_shape
-    expected_output_shape[pad_axis] = instances
-    output = preprocess_ops.pad_max_instances(
-        np.ones(input_shape), instances, pad_axis=pad_axis)
-    self.assertAllEqual(expected_output_shape, tf.shape(output).numpy())
-
-
-if __name__ == '__main__':
-  tf.test.main()

official/vision/beta/projects/yolo/ops/preprocessing_ops_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.
+
+"""Tests for preprocessing_ops.py."""
+from absl.testing import parameterized
+import numpy as np
+import tensorflow as tf
+
+from official.vision.beta.ops import box_ops as bbox_ops
+from official.vision.beta.projects.yolo.ops import preprocessing_ops
+
+
+class InputUtilsTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.parameters(([1, 2], 20, 0), ([13, 2, 4], 15, 0))
+  def testPadMaxInstances(self, input_shape, instances, pad_axis):
+    expected_output_shape = input_shape
+    expected_output_shape[pad_axis] = instances
+    output = preprocessing_ops.pad_max_instances(
+        np.ones(input_shape), instances, pad_axis=pad_axis)
+    self.assertAllEqual(expected_output_shape, tf.shape(output).numpy())
+
+  @parameterized.parameters((100, 200))
+  def testGetImageShape(self, image_height, image_width):
+    image = tf.convert_to_tensor(np.random.rand(image_height, image_width, 3))
+    image_shape = preprocessing_ops.get_image_shape(image)
+    self.assertAllEqual((image_height, image_width), image_shape)
+
+  @parameterized.parameters((400, 600, .5, .5, .0, True),
+                            (100, 200, .5, .5, .5))
+  def testImageRandHSV(self,
+                       image_height,
+                       image_width,
+                       rh,
+                       rs,
+                       rv,
+                       is_darknet=False):
+    image = tf.convert_to_tensor(np.random.rand(image_height, image_width, 3))
+    processed_image = preprocessing_ops.image_rand_hsv(
+        image, rh, rs, rv, darknet=is_darknet)
+    processed_image_shape = tf.shape(processed_image)
+    self.assertAllEqual([image_height, image_width, 3],
+                        processed_image_shape.numpy())
+
+  @parameterized.parameters((100, 200, [50, 100]))
+  def testResizeAndJitterImage(self, image_height, image_width, desired_size):
+    image = tf.convert_to_tensor(np.random.rand(image_height, image_width, 3))
+    processed_image, _, _ = preprocessing_ops.resize_and_jitter_image(
+        image, desired_size)
+    processed_image_shape = tf.shape(processed_image)
+    self.assertAllEqual([desired_size[0], desired_size[1], 3],
+                        processed_image_shape.numpy())
+
+  @parameterized.parameters((400, 600, [200, 300]))
+  def testAffineWarpImage(self,
+                          image_height,
+                          image_width,
+                          desired_size,
+                          degrees=7.0,
+                          scale_min=0.1,
+                          scale_max=1.9):
+    image = tf.convert_to_tensor(np.random.rand(image_height, image_width, 3))
+    processed_image, _, _ = preprocessing_ops.affine_warp_image(
+        image,
+        desired_size,
+        degrees=degrees,
+        scale_min=scale_min,
+        scale_max=scale_max)
+    processed_image_shape = tf.shape(processed_image)
+    self.assertAllEqual([desired_size[0], desired_size[1], 3],
+                        processed_image_shape.numpy())
+
+  # Working Test
+  @parameterized.parameters(([[400, 600], [200, 300],
+                              [[0, 0, 0], [0, 0, 0], [0, 0, 0]]], 50))
+  def testAffineWarpBoxes(self, affine, num_boxes):
+    boxes = tf.convert_to_tensor(np.random.rand(num_boxes, 4))
+    boxes = bbox_ops.denormalize_boxes(boxes, affine[0])
+    processed_boxes, _ = preprocessing_ops.affine_warp_boxes(
+        tf.cast(affine[2], tf.double), boxes, affine[1], box_history=boxes)
+    processed_boxes_shape = tf.shape(processed_boxes)
+    self.assertAllEqual([num_boxes, 4], processed_boxes_shape.numpy())
+
+  # Working Test
+  @parameterized.parameters(([100, 100], [[-0.489, 51.28, 0.236, 51.686],
+                                          [65, 100, 200, 150],
+                                          [150, 80, 200, 130]]))
+  def testBoxCandidates(self, output_size, boxes):
+    boxes = tf.cast(bbox_ops.denormalize_boxes(boxes, output_size), tf.double)
+    clipped_ind = preprocessing_ops.boxes_candidates(
+        boxes, boxes, ar_thr=1e32, wh_thr=0, area_thr=tf.cast(0, tf.double))
+    clipped_ind_shape = tf.shape(clipped_ind)
+    self.assertAllEqual([3], clipped_ind_shape.numpy())
+    self.assertAllEqual([0, 1, 2], clipped_ind.numpy())
+
+  # Working Test
+  @parameterized.parameters((
+      50,
+      [0.5, 0.5],
+      [0, 0],  # Clipping all boxes
+      [0.0, 0.0]))
+  def testResizeAndCropBoxes(self, num_boxes, image_scale, output_size, offset):
+    boxes = tf.convert_to_tensor(np.random.rand(num_boxes, 4))
+    processed_boxes, _ = preprocessing_ops.resize_and_crop_boxes(
+        boxes, tf.cast(image_scale, tf.double), output_size,
+        tf.cast(offset, tf.double), boxes)
+    processed_boxes_shape = tf.shape(processed_boxes)
+    self.assertAllEqual([num_boxes, 4], processed_boxes_shape.numpy())
+    self.assertAllEqual(
+        tf.math.reduce_sum(processed_boxes), tf.convert_to_tensor(0))
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/vision/beta/projects/yolo/tasks/image_classification.py

@@ -12,15 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Lint as: python3
 """Image classification task definition."""
-import tensorflow as tf
-
-from official.core import input_reader
+from official.common import dataset_fn
 from official.core import task_factory
-from official.vision.beta.dataloaders import classification_input
+from official.vision.beta.dataloaders import classification_input as classification_input_base
+from official.vision.beta.dataloaders import input_reader_factory
+from official.vision.beta.dataloaders import tfds_factory
 from official.vision.beta.projects.yolo.configs import darknet_classification as exp_cfg
-from official.vision.beta.projects.yolo.dataloaders import classification_tfds_decoder as cli
+from official.vision.beta.projects.yolo.dataloaders import classification_input
 from official.vision.beta.tasks import image_classification
 
 
@@ -33,82 +32,34 @@ class ImageClassificationTask(image_classification.ImageClassificationTask):
 
     num_classes = self.task_config.model.num_classes
     input_size = self.task_config.model.input_size
+    image_field_key = self.task_config.train_data.image_field_key
+    label_field_key = self.task_config.train_data.label_field_key
+    is_multilabel = self.task_config.train_data.is_multilabel
 
     if params.tfds_name:
-      decoder = cli.Decoder()
+      decoder = tfds_factory.get_classification_decoder(params.tfds_name)
     else:
-      decoder = classification_input.Decoder()
+      decoder = classification_input_base.Decoder(
+          image_field_key=image_field_key,
+          label_field_key=label_field_key,
+          is_multilabel=is_multilabel)
 
     parser = classification_input.Parser(
         output_size=input_size[:2],
         num_classes=num_classes,
+        image_field_key=image_field_key,
+        label_field_key=label_field_key,
+        decode_jpeg_only=params.decode_jpeg_only,
+        aug_rand_hflip=params.aug_rand_hflip,
+        aug_type=params.aug_type,
+        is_multilabel=is_multilabel,
         dtype=params.dtype)
 
-    reader = input_reader.InputReader(
+    reader = input_reader_factory.input_reader_generator(
         params,
-        dataset_fn=tf.data.TFRecordDataset,
+        dataset_fn=dataset_fn.pick_dataset_fn(params.file_type),
         decoder_fn=decoder.decode,
         parser_fn=parser.parse_fn(params.is_training))
 
     dataset = reader.read(input_context=input_context)
     return dataset
-
-  def train_step(self, inputs, model, optimizer, metrics=None):
-    """Does forward and backward.
-
-    Args:
-      inputs: a dictionary of input tensors.
-      model: the model, forward pass definition.
-      optimizer: the optimizer for this training step.
-      metrics: a nested structure of metrics objects.
-
-    Returns:
-      A dictionary of logs.
-    """
-    features, labels = inputs
-    if self.task_config.losses.one_hot:
-      labels = tf.one_hot(labels, self.task_config.model.num_classes)
-
-    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
-    with tf.GradientTape() as tape:
-      outputs = model(features, training=True)
-      # Casting output layer as float32 is necessary when mixed_precision is
-      # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
-      outputs = tf.nest.map_structure(
-          lambda x: tf.cast(x, tf.float32), outputs)
-
-      # Computes per-replica loss.
-      loss = self.build_losses(
-          model_outputs=outputs, labels=labels, aux_losses=model.losses)
-      # Scales loss as the default gradients allreduce performs sum inside the
-      # optimizer.
-      scaled_loss = loss / num_replicas
-
-      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
-      # scaled for numerical stability.
-      if isinstance(optimizer, tf.keras.mixed_precision.LossScaleOptimizer):
-        scaled_loss = optimizer.get_scaled_loss(scaled_loss)
-
-    tvars = model.trainable_variables
-    grads = tape.gradient(scaled_loss, tvars)
-    # Scales back gradient before apply_gradients when LossScaleOptimizer is
-    # used.
-    if isinstance(optimizer, tf.keras.mixed_precision.LossScaleOptimizer):
-      grads = optimizer.get_unscaled_gradients(grads)
-
-    # Apply gradient clipping.
-    if self.task_config.gradient_clip_norm > 0:
-      grads, _ = tf.clip_by_global_norm(
-          grads, self.task_config.gradient_clip_norm)
-    optimizer.apply_gradients(list(zip(grads, tvars)))
-
-    logs = {self.loss: loss}
-    if metrics:
-      self.process_metrics(metrics, labels, outputs)
-      logs.update({m.name: m.result() for m in metrics})
-    elif model.compiled_metrics:
-      self.process_compiled_metrics(model.compiled_metrics, labels, outputs)
-      logs.update({m.name: m.result() for m in model.metrics})
-    return logs
-
-