Internal change

PiperOrigin-RevId: 358867611

official/vision/beta/dataloaders/input_reader_factory.py

+# Lint as: python3
+# Copyright 2020 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.
+# ==============================================================================
+"""Factory for getting TF-Vision input readers."""
+
+from official.common import dataset_fn as dataset_fn_util
+from official.core import config_definitions as cfg
+from official.core import input_reader as core_input_reader
+
+from official.vision.beta.dataloaders import input_reader as vision_input_reader
+
+
+def input_reader_generator(params: cfg.DataConfig,
+                           **kwargs) -> core_input_reader.InputReader:
+  """Instantiates an input reader class according to the params.
+
+  Args:
+    params: A config_definitions.DataConfig object.
+    **kwargs: Additional arguments passed to input reader initialization.
+
+  Returns:
+    An InputReader object.
+
+  """
+  if params.is_training and params.get('pseudo_label_data', False):
+    return vision_input_reader.CombinationDatasetInputReader(
+        params,
+        pseudo_label_dataset_fn=dataset_fn_util.pick_dataset_fn(
+            params.pseudo_label_data.file_type),
+        **kwargs)
+  else:
+    return core_input_reader.InputReader(params, **kwargs)

official/vision/beta/tasks/image_classification.py

@@ -16,13 +16,14 @@
 """Image classification task definition."""
 from absl import logging
 import tensorflow as tf
+
 from official.common import dataset_fn
 from official.core import base_task
-from official.core import input_reader
 from official.core import task_factory
 from official.modeling import tf_utils
 from official.vision.beta.configs import image_classification as exp_cfg
 from official.vision.beta.dataloaders import classification_input
+from official.vision.beta.dataloaders import input_reader_factory
 from official.vision.beta.dataloaders import tfds_classification_decoders
 from official.vision.beta.modeling import factory
 
@@ -96,7 +97,7 @@ class ImageClassificationTask(base_task.Task):
         randaug_magnitude=params.randaug_magnitude,
         dtype=params.dtype)
 
-    reader = input_reader.InputReader(
+    reader = input_reader_factory.input_reader_generator(
         params,
         dataset_fn=dataset_fn.pick_dataset_fn(params.file_type),
         decoder_fn=decoder.decode,

official/vision/beta/tasks/maskrcnn.py

@@ -19,9 +19,9 @@ from absl import logging
 import tensorflow as tf
 from official.common import dataset_fn
 from official.core import base_task
-from official.core import input_reader
 from official.core import task_factory
 from official.vision.beta.configs import maskrcnn as exp_cfg
+from official.vision.beta.dataloaders import input_reader_factory
 from official.vision.beta.dataloaders import maskrcnn_input
 from official.vision.beta.dataloaders import tf_example_decoder
 from official.vision.beta.dataloaders import tf_example_label_map_decoder
@@ -143,7 +143,7 @@ class MaskRCNNTask(base_task.Task):
         include_mask=self._task_config.model.include_mask,
         mask_crop_size=params.parser.mask_crop_size)
 
-    reader = input_reader.InputReader(
+    reader = input_reader_factory.input_reader_generator(
         params,
         dataset_fn=dataset_fn.pick_dataset_fn(params.file_type),
         decoder_fn=decoder.decode,

official/vision/beta/tasks/retinanet.py

@@ -19,10 +19,10 @@ from absl import logging
 import tensorflow as tf
 from official.common import dataset_fn
 from official.core import base_task
-from official.core import input_reader
 from official.core import task_factory
 from official.vision import keras_cv
 from official.vision.beta.configs import retinanet as exp_cfg
+from official.vision.beta.dataloaders import input_reader_factory
 from official.vision.beta.dataloaders import retinanet_input
 from official.vision.beta.dataloaders import tf_example_decoder
 from official.vision.beta.dataloaders import tfds_detection_decoders
@@ -122,7 +122,7 @@ class RetinaNetTask(base_task.Task):
         skip_crowd_during_training=params.parser.skip_crowd_during_training,
         max_num_instances=params.parser.max_num_instances)
 
-    reader = input_reader.InputReader(
+    reader = input_reader_factory.input_reader_generator(
         params,
         dataset_fn=dataset_fn.pick_dataset_fn(params.file_type),
         decoder_fn=decoder.decode,

official/vision/beta/tasks/semantic_segmentation.py

@@ -19,9 +19,9 @@ from absl import logging
 import tensorflow as tf
 from official.common import dataset_fn
 from official.core import base_task
-from official.core import input_reader
 from official.core import task_factory
 from official.vision.beta.configs import semantic_segmentation as exp_cfg
+from official.vision.beta.dataloaders import input_reader_factory
 from official.vision.beta.dataloaders import segmentation_input
 from official.vision.beta.dataloaders import tfds_segmentation_decoders
 from official.vision.beta.evaluation import segmentation_metrics
@@ -104,7 +104,7 @@ class SemanticSegmentationTask(base_task.Task):
         aug_rand_hflip=params.aug_rand_hflip,
         dtype=params.dtype)
 
-    reader = input_reader.InputReader(
+    reader = input_reader_factory.input_reader_generator(
         params,
         dataset_fn=dataset_fn.pick_dataset_fn(params.file_type),
         decoder_fn=decoder.decode,

official/vision/beta/tasks/video_classification.py

@@ -17,10 +17,10 @@
 from absl import logging
 import tensorflow as tf
 from official.core import base_task
-from official.core import input_reader
 from official.core import task_factory
 from official.modeling import tf_utils
 from official.vision.beta.configs import video_classification as exp_cfg
+from official.vision.beta.dataloaders import input_reader_factory
 from official.vision.beta.dataloaders import video_input
 from official.vision.beta.modeling import factory_3d
 
@@ -74,7 +74,7 @@ class VideoClassificationTask(base_task.Task):
     parser = video_input.Parser(input_params=params)
     postprocess_fn = video_input.PostBatchProcessor(params)
 
-    reader = input_reader.InputReader(
+    reader = input_reader_factory.input_reader_generator(
         params,
         dataset_fn=self._get_dataset_fn(params),
         decoder_fn=self._get_decoder_fn(params),