Fix usage of optional dependencies (#12)

Author: fcogidi

mmlearn/conf/__init__.py

@@ -1,6 +1,7 @@
 """Hydra/Hydra-zen-based configurations."""
 
 import functools
+import warnings
 from dataclasses import dataclass, field
 from enum import Enum
 from pathlib import Path
@@ -446,6 +447,11 @@ class name is used.
         group="trainer/logger",
         provider="lightning",
     )
+else:
+    warnings.warn(
+        "wandb is not available. Skipping registration of 'trainer/logger/WandbLogger'.",
+        stacklevel=1,
+    )
 
 
 #################### Custom Hydra Main Decorator ####################

mmlearn/datasets/librispeech.py

@@ -5,14 +5,15 @@
 import torch
 import torch.nn.functional as F  # noqa: N812
 from hydra_zen import MISSING, store
+from lightning_utilities.core.imports import RequirementCache
 from torch.utils.data.dataset import Dataset
-from torchaudio.datasets import LIBRISPEECH
 
 from mmlearn.constants import EXAMPLE_INDEX_KEY
 from mmlearn.datasets.core import Modalities
 from mmlearn.datasets.core.example import Example
 
 
+_TORCHAUDIO_AVAILABLE = RequirementCache("torchaudio>=2.4.0")
 SAMPLE_RATE = 16000
 
 
@@ -80,6 +81,12 @@ class LibriSpeech(Dataset[Example]):
     def __init__(self, root_dir: str, split: str = "train-clean-100") -> None:
         """Initialize LibriSpeech dataset."""
         super().__init__()
+        if not _TORCHAUDIO_AVAILABLE:
+            raise ImportError(
+                "LibriSpeech dataset requires `torchaudio` which is not installed."
+            )
+        from torchaudio.datasets import LIBRISPEECH
+
         self.dataset = LIBRISPEECH(
             root=root_dir,
             url=split,

mmlearn/datasets/nyuv2.py

@@ -3,11 +3,10 @@
 import os
 from typing import Callable, List, Literal, Optional
 
-import cv2
 import numpy as np
 import torch
 from hydra_zen import MISSING, store
-from PIL import Image
+from lightning_utilities.core.imports import RequirementCache
 from PIL.Image import Image as PILImage
 from torch.utils.data import Dataset
 from torchvision.transforms.v2.functional import to_pil_image
@@ -17,6 +16,11 @@
 from mmlearn.datasets.core.example import Example
 
 
+_OPENCV_AVAILABLE = RequirementCache("opencv-python>=4.10.0.84")
+if _OPENCV_AVAILABLE:
+    import cv2  # noqa: F401
+
+
 _LABELS = [
     "bedroom",
     "kitchen",
@@ -56,7 +60,7 @@ def depth_normalize(
     torch.Tensor
         Normalized depth image.
     """
-    depth_image = np.array(Image.open(depth_file))
+    depth_image = np.array(PILImage.open(depth_file))
     depth = np.array(depth_image).astype(np.float32)
     depth_in_meters = depth / 1000.0
 
@@ -100,6 +104,10 @@ def __init__(
         depth_transform: Optional[Callable[[PILImage], torch.Tensor]] = None,
     ) -> None:
         super().__init__()
+        if not _OPENCV_AVAILABLE:
+            raise ImportError(
+                "NYUv2 dataset requires `opencv-python` which is not installed.",
+            )
         self._validate_args(root_dir, split, rgb_transform, depth_transform)
         self.return_type = return_type
 

mmlearn/datasets/sunrgbd.py

@@ -3,11 +3,10 @@
 import os
 from typing import Callable, List, Literal, Optional
 
-import cv2
 import numpy as np
 import torch
 from hydra_zen import MISSING, store
-from PIL import Image
+from lightning_utilities.core.imports import RequirementCache
 from PIL.Image import Image as PILImage
 from torch.utils.data import Dataset
 from torchvision.transforms.v2.functional import to_pil_image
@@ -17,6 +16,10 @@
 from mmlearn.datasets.core.example import Example
 
 
+_OPENCV_AVAILABLE = RequirementCache("opencv-python>=4.10.0.84")
+if _OPENCV_AVAILABLE:
+    import cv2  # noqa: F401
+
 _LABELS = [
     "bathroom",
     "bedroom",
@@ -97,7 +100,7 @@ def convert_depth_to_disparity(
         lines = fh.readlines()
         focal_length = float(lines[0].strip().split()[0])
     baseline = sensor_to_params[sensor_type]["baseline"]
-    depth_image = np.array(Image.open(depth_file))
+    depth_image = np.array(PILImage.open(depth_file))
     depth = np.array(depth_image).astype(np.float32)
     depth_in_meters = depth / 1000.0
     if min_depth is not None:
@@ -143,6 +146,11 @@ def __init__(
         depth_transform: Optional[Callable[[PILImage], torch.Tensor]] = None,
     ) -> None:
         super().__init__()
+        if not _OPENCV_AVAILABLE:
+            raise ImportError(
+                "SUN RGB-D dataset requires `opencv-python` which is not installed.",
+            )
+
         self._validate_args(root_dir, split, rgb_transform, depth_transform)
         self.return_type = return_type