Smarter optional dependency imports in the datasets. (OpenNMT#1195)

onmt/inputters/audio_dataset.py

@@ -6,8 +6,21 @@
 
 from onmt.inputters.dataset_base import DatasetBase
 
+# imports of datatype-specific dependencies
+try:
+    import torchaudio
+    import librosa
+    import numpy as np
+except ImportError:
+    torchaudio, librosa, np = None, None, None
+
 
 class AudioDataset(DatasetBase):
+    @staticmethod
+    def _check_deps():
+        if any([torchaudio is None, librosa is None, np is None]):
+            AudioDataset._raise_missing_dep(
+                "torchaudio", "librosa", "numpy")
 
     @staticmethod
     def sort_key(ex):
@@ -17,13 +30,11 @@ def sort_key(ex):
     @staticmethod
     def extract_features(audio_path, sample_rate, truncate, window_size,
                          window_stride, window, normalize_audio):
-        import torchaudio
-        import librosa
-        import numpy as np
         # torchaudio loading options recently changed. It's probably
         # straightforward to rewrite the audio handling to make use of
         # up-to-date torchaudio, but in the meantime there is a legacy
         # method which uses the old defaults
+        AudioDataset._check_deps()
         sound, sample_rate_ = torchaudio.legacy.load(audio_path)
         if truncate and truncate > 0:
             if sound.size(0) > truncate:

onmt/inputters/dataset_base.py

@@ -9,6 +9,12 @@
 from torchtext.vocab import Vocab
 
 
+# several data readers need optional dependencies. There's no
+# appropriate builtin exception
+class MissingDependencyException(Exception):
+    pass
+
+
 class DatasetBase(Dataset):
     """
     A dataset is an object that accepts sequences of raw data (sentence pairs
@@ -80,6 +86,13 @@ def __init__(self, fields, src_examples_iter, tgt_examples_iter,
 
         super(DatasetBase, self).__init__(examples, fields, filter_pred)
 
+    @staticmethod
+    def _raise_missing_dep(*missing_deps):
+        """Raise missing dep exception with standard error message."""
+        raise MissingDependencyException(
+            "Could not create reader. Be sure to install "
+            "the following dependencies: " + ", ".join(missing_deps))
+
     def __getattr__(self, attr):
         # avoid infinite recursion when fields isn't defined
         if 'fields' not in vars(self):

onmt/inputters/image_dataset.py

@@ -4,8 +4,21 @@
 
 from onmt.inputters.dataset_base import DatasetBase
 
+# domain specific dependencies
+try:
+    from PIL import Image
+    from torchvision import transforms
+    import cv2
+except ImportError:
+    Image, transforms, cv2 = None, None, None
+
 
 class ImageDataset(DatasetBase):
+    @staticmethod
+    def _check_deps():
+        if any([Image is None, transforms is None, cv2 is None]):
+            ImageDataset._raise_missing_dep(
+                "PIL", "torchvision", "cv2")
 
     @staticmethod
     def sort_key(ex):
@@ -25,9 +38,7 @@ def make_examples(
         Yields:
             a dictionary containing image data, path and index for each line.
         """
-        from PIL import Image
-        from torchvision import transforms
-        import cv2
+        ImageDataset._check_deps()
 
         if isinstance(images, str):
             images = cls._read_file(images)