add missing colorize dataset

Summary: Pull Request resolved: https://github.com/fairinternal/fairseq-py/pull/851

Differential Revision: D17145769

Pulled By: alexeib

fbshipit-source-id: 9dd26799d044ae5386e8204a129b5e3fc66d6e85

fairseq/data/__init__.py

@@ -11,6 +11,7 @@ from .base_wrapper_dataset import BaseWrapperDataset
 
 from .audio.raw_audio_dataset import FileAudioDataset
 from .backtranslation_dataset import BacktranslationDataset
+from .colorize_dataset import ColorizeDataset
 from .concat_dataset import ConcatDataset
 from .concat_sentences_dataset import ConcatSentencesDataset
 from .id_dataset import IdDataset
@@ -51,6 +52,7 @@ from .iterators import (
 __all__ = [
     'BacktranslationDataset',
     'BaseWrapperDataset',
+    'ColorizeDataset',
     'ConcatDataset',
     'ConcatSentencesDataset',
     'CountingIterator',

fairseq/data/colorize_dataset.py

+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+
+from . import BaseWrapperDataset
+
+
+class ColorizeDataset(BaseWrapperDataset):
+    """ Adds 'colors' property to net input that is obtained from the provided color getter for use by models """
+    def __init__(self, dataset, color_getter):
+        super().__init__(dataset)
+        self.color_getter = color_getter
+
+    def collater(self, samples):
+        base_collate = super().collater(samples)
+        if len(base_collate) > 0:
+            base_collate["net_input"]["colors"] = torch.tensor(
+                list(self.color_getter(self.dataset, s["id"]) for s in samples),
+                dtype=torch.long,
+            )
+        return base_collate