Introduce code from Rust monorepo. This is a mess

tts/nemo/collections/cv/datasets/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  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.
+
+from nemo.collections.cv.datasets.mnist_dataset import MNISTDataset, MNISTDatasetConfig

tts/nemo/collections/cv/datasets/mnist_dataset.py

@@ -0,0 +1,134 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2020, NVIDIA CORPORATION.  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.
+
+from dataclasses import dataclass
+from os.path import expanduser
+from typing import Optional
+
+from torchvision.datasets import MNIST
+from torchvision.transforms import Compose, Resize, ToTensor
+
+from nemo.core.classes import Dataset
+from nemo.core.config import Config
+from nemo.core.neural_types.axes import AxisKind, AxisType
+from nemo.core.neural_types.elements import ClassificationTarget, Index, NormalizedImageValue, StringLabel
+from nemo.core.neural_types.neural_type import NeuralType
+from nemo.utils.decorators import add_port_docs, experimental
+
+
+@dataclass
+class MNISTDatasetConfig(Config):
+    """
+    Structured config for MNISTDataset class.
+
+    Args:
+        height: image height (DEFAULT: 28)
+        width: image width (DEFAULT: 28)
+        data_folder: path to the folder with data, can be relative to user (DEFAULT: "~/data/mnist")
+        train: use train or test splits (DEFAULT: True)
+        name: Name of the module (DEFAULT: None)
+    """
+
+    height: int = 28
+    width: int = 28
+    data_folder: str = "~/data/mnist"
+    train: bool = True
+    download: bool = True
+
+
+class MNISTDataset(Dataset):
+    """
+    A "thin wrapper" around the torchvision's MNIST dataset.
+    """
+
+    def __init__(self, cfg: MNISTDatasetConfig = MNISTDatasetConfig()):
+        """
+        Initializes the MNIST dataset.
+
+        Args:
+            cfg: Configuration object of type MNISTDatasetConfig.
+
+        """
+        # Call the base class constructor of Dataset.
+        Dataset.__init__(self)  # , name=name)
+
+        # Store height and width.
+        self._height = cfg.height
+        self._width = cfg.width
+
+        # Create transformations: up-scale and transform to tensors.
+        mnist_transforms = Compose([Resize((self._height, self._width)), ToTensor()])
+
+        # Get absolute path.
+        abs_data_folder = expanduser(cfg.data_folder)
+
+        # Create the MNIST dataset object.
+        self._dataset = MNIST(root=abs_data_folder, train=cfg.train, download=cfg.download, transform=mnist_transforms)
+
+        # Class names.
+        labels = 'Zero One Two Three Four Five Six Seven Eight Nine'.split(' ')
+        word_to_ix = {labels[i]: i for i in range(10)}
+
+        # Reverse mapping.
+        self._ix_to_word = {value: key for (key, value) in word_to_ix.items()}
+
+    @property
+    def output_types(self):
+        """
+        Creates definitions of output ports.
+        """
+        return {
+            "indices": NeuralType(tuple('B'), elements_type=Index()),
+            "images": NeuralType(
+                axes=(
+                    AxisType(kind=AxisKind.Batch),
+                    AxisType(kind=AxisKind.Channel, size=1),
+                    AxisType(kind=AxisKind.Height, size=self._height),
+                    AxisType(kind=AxisKind.Width, size=self._width),
+                ),
+                elements_type=NormalizedImageValue(),  # float, <0-1>
+            ),
+            "targets": NeuralType(tuple('B'), elements_type=ClassificationTarget()),  # Target are ints!
+            "labels": NeuralType(tuple('B'), elements_type=StringLabel()),  # Labels is string!
+        }
+
+    def __len__(self):
+        """
+        Returns:
+            Length of the dataset.
+        """
+        return len(self._dataset)
+
+    def __getitem__(self, index: int):
+        """
+        Returns a single sample.
+
+        Args:
+            index: index of the sample to return.
+        """
+        # Get image and target.
+        img, target = self._dataset.__getitem__(index)
+
+        # Return sample.
+        return index, img, target, self._ix_to_word[target]
+
+    @property
+    def ix_to_word(self):
+        """
+        Returns:
+            Dictionary with mapping of target indices (int) to labels (class names as strings)
+            that can we used by other modules.
+        """
+        return self._ix_to_word