fix(datamodules): catch ModuleNotFound datamodules

Author: okunator

cellseg_models_pytorch/__init__.py

@@ -1,17 +1,14 @@
-from . import datamodules, datasets, inference, models, training, utils
+from . import inference, models, utils
 from .models import CellPoseUnet, HoverNet, StarDistUnet
 
 __version__ = "0.1.0"
-submodules = ["utils", "models", "inference", "training", "datasets", "datamodules"]
+submodules = ["utils", "models", "inference"]
 __all__ = [
     "__version__",
-    "utils",
     "CellPoseUnet",
     "StarDistUnet",
     "HoverNet",
+    "utils",
     "models",
     "inference",
-    "training",
-    "datasets",
-    "datamodules",
 ]

cellseg_models_pytorch/datamodules/_basemodule.py

@@ -2,7 +2,13 @@
 from pathlib import Path
 from typing import Union
 
-import pytorch_lightning as pl
+try:
+    import pytorch_lightning as pl
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "To use the `csmp.datamodules` module, the pytorch-lightning lib, is needed. "
+        "Install with `pip install pytorch-lightning`"
+    )
 from torch.utils.data import DataLoader
 
 __all__ = ["BaseDataModule"]

cellseg_models_pytorch/datamodules/downloader.py

@@ -1,7 +1,13 @@
 import re
 from pathlib import Path
 
-import requests
+try:
+    import requests
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "To use `SimpleDownloader`, the requests package is needed. "
+        "Install with `pip install requests`."
+    )
 from tqdm import tqdm
 
 

cellseg_models_pytorch/datamodules/pannuke_datamodule.py

@@ -5,10 +5,18 @@
 import numpy as np
 from tqdm import tqdm
 
-from ..datasets import SegmentationFolderDataset
 from ..utils import FileHandler, fix_duplicates
-from ._basemodule import BaseDataModule
-from .downloader import SimpleDownloader
+
+try:
+    from ..datasets import SegmentationFolderDataset
+    from ._basemodule import BaseDataModule
+    from .downloader import SimpleDownloader
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "To use the PannukeDataModule, requests, pytorch-lightning, & albumentations "
+        "libraries are needed. Install with "
+        "`pip install requests pytorch-lightning albumentations`"
+    )
 
 
 class PannukeDataModule(BaseDataModule):