Setup instructions, docstrings

README.md

@@ -1 +1,45 @@
-# diabetic-retinopathy-detection
+# Diabetic Retinopathy Detection with AI
+
+## Setup
+
+### Gradio app environment
+
+Install from pip requirements file:
+
+```bash
+conda create -n retinopathy_app python=3.10
+conda activate retinopathy_app
+pip install -r requirements.txt
+python app.py
+```
+
+Install manually:
+
+```bash
+pip install pytorch --index-url  https://download.pytorch.org/whl/cpu
+pip install gradio
+pip install transformers
+```
+
+### Training environment
+
+Create conda environment from YAML:
+```bash
+mamba env create -n retinopathy_train -f environment.yml
+```
+
+Download the data from [Kaggle](https://www.kaggle.com/competitions/diabetic-retinopathy-detection/data) or use kaggle API:
+
+```bash
+pip install kaggle
+kaggle competitions download -c diabetic-retinopathy-detection
+mkdir retinopathy_data/
+unzip diabetic-retinopathy-detection.zip -d retinopathy_data/
+```
+
+Launch training:
+```bash
+conda activate retinopathy_train
+python train.py
+```
+The trained model will be put into `lightning_logs/`.

app.py

@@ -2,7 +2,7 @@ import os
 import gradio as gr
 import numpy as np
 import torch
-from typing import Tuple, Optional, Dict, List
+from typing import Tuple, Optional, Dict, List, Dict
 import glob
 from collections import defaultdict
 
@@ -13,7 +13,10 @@ from labelmap import DR_LABELMAP
 
 
 class App:
+    """ Demonstration of the Diabetic Retinopathy model as a Gradio app. """
+
     def __init__(self) -> None:
+        """ Constructor. """
 
         ckpt_name = "2023-12-24_20-02-18_30345221_V100_x4_resnet34/"
 
@@ -41,7 +44,7 @@ class App:
                     output = gr.Label(num_top_classes=len(DR_LABELMAP),
                                       label="Retinopathy level prediction")
                 with gr.Column(scale=4):
-                    gr.Markdown("![](https://media.githubusercontent.com/media/Obs01ete/retinopathy/master/media/logo1.png)")
+                    gr.Markdown("![](https://media.githubusercontent.com/media/SDAIA-KAUST-AI/diabetic-retinopathy-detection/main/media/logo1.png)")
             with gr.Row():
                 with gr.Column(scale=9, min_width=100):
                     image = gr.Image(label="Retina scan")
@@ -66,9 +69,19 @@ class App:
         self.ui = ui
 
     def launch(self) -> None:
+        """ Launch the application, blocking. """
         self.ui.queue().launch(share=True)
 
-    def predict(self, image: Optional[np.ndarray]):
+    def predict(self, image: Optional[np.ndarray]) -> Dict[str, float]:
+        """ Gradio callback for pricessing of an image.
+
+        Args:
+            image (Optional[np.ndarray]): Provided image.
+
+        Returns:
+            Dict[str, float]: Label-compatible dict.
+        """
+
         if image is None:
             return dict()
         cls_name, prob, probs = self._infer(image)
@@ -79,6 +92,19 @@ class App:
         return probs_dict
     
     def _infer(self, image_chw: np.ndarray) -> Tuple[str, float, np.ndarray]:
+        """ Low-level method to perform neural network inference.
+
+        Args:
+            image_chw (np.ndarray): Provided image.
+
+        Returns:
+            Tuple[str, float, np.ndarray]:
+                - Most probable class name
+                - Probability of the most probable class name.
+                - Probablilities of all classes in the order of
+                  being listed in the label map.
+        """
+
         assert isinstance(self.model, ResNetForImageClassification)
 
         inputs = self.image_processor(image_chw, return_tensors="pt")
@@ -98,6 +124,11 @@ class App:
 
     @staticmethod
     def _load_example_lists() -> Dict[int, List[str]]:
+        """ Load example retina images from disk.
+
+        Returns:
+            Dict[int, List[str]]: Dictionary of cls_id -> list of images paths.
+        """
 
         example_flat_list = glob.glob("demo_data/train/**/*.jpeg")
 
@@ -115,6 +146,7 @@ class App:
 
 
 def main():
+    """ App entry point. """
     app = App()
     app.launch()
 

environment.yml

@@ -1,22 +1,251 @@
-name: diabetic-retinopathy-detection
-
+name: retinopathy
 channels:
-  - pytorch
-  - nvidia
+  - anaconda
   - conda-forge
   - defaults
-
 dependencies:
-  - dask
-  - gh
-  - git
-  - kaggle
-  - jupyterlab
-  - jupyterlab-nvdashboard
-  - lightning
-  - nbgitpuller
-  - pytorch
-  - pytorch-cuda=11.8
-  - scikit-learn
-  - torchvision
-  
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=5.1=1_gnu
+  - aiofiles=22.1.0=py310h06a4308_0
+  - aiosqlite=0.18.0=py310h06a4308_0
+  - argon2-cffi=21.3.0=pyhd3eb1b0_0
+  - argon2-cffi-bindings=21.2.0=py310h7f8727e_0
+  - asttokens=2.0.5=pyhd3eb1b0_0
+  - attrs=23.1.0=py310h06a4308_0
+  - babel=2.11.0=py310h06a4308_0
+  - backcall=0.2.0=pyhd3eb1b0_0
+  - beautifulsoup4=4.12.2=py310h06a4308_0
+  - bleach=4.1.0=pyhd3eb1b0_0
+  - brotli-python=1.0.9=py310h6a678d5_7
+  - bzip2=1.0.8=h7b6447c_0
+  - ca-certificates=2023.08.22=h06a4308_0
+  - certifi=2023.11.17=py310h06a4308_0
+  - cffi=1.16.0=py310h5eee18b_0
+  - comm=0.1.2=py310h06a4308_0
+  - cryptography=41.0.3=py310hdda0065_0
+  - debugpy=1.6.7=py310h6a678d5_0
+  - decorator=5.1.1=pyhd3eb1b0_0
+  - defusedxml=0.7.1=pyhd3eb1b0_0
+  - executing=0.8.3=pyhd3eb1b0_0
+  - ipykernel=6.25.0=py310h2f386ee_0
+  - ipython=8.15.0=py310h06a4308_0
+  - ipython_genutils=0.2.0=pyhd3eb1b0_1
+  - jedi=0.18.1=py310h06a4308_1
+  - jinja2=3.1.2=py310h06a4308_0
+  - json5=0.9.6=pyhd3eb1b0_0
+  - jsonschema=4.19.2=py310h06a4308_0
+  - jsonschema-specifications=2023.7.1=py310h06a4308_0
+  - jupyter_client=8.6.0=py310h06a4308_0
+  - jupyter_core=5.5.0=py310h06a4308_0
+  - jupyter_events=0.8.0=py310h06a4308_0
+  - jupyter_server=2.10.0=py310h06a4308_0
+  - jupyter_server_fileid=0.9.0=py310h06a4308_0
+  - jupyter_server_terminals=0.4.4=py310h06a4308_1
+  - jupyter_server_ydoc=0.8.0=py310h06a4308_1
+  - jupyter_ydoc=0.2.4=py310h06a4308_0
+  - jupyterlab=3.6.3=py310h06a4308_0
+  - jupyterlab_pygments=0.2.2=py310h06a4308_0
+  - jupyterlab_server=2.25.1=py310h06a4308_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - libffi=3.4.4=h6a678d5_0
+  - libgcc-ng=11.2.0=h1234567_1
+  - libgomp=11.2.0=h1234567_1
+  - libsodium=1.0.18=h7b6447c_0
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - libuuid=1.41.5=h5eee18b_0
+  - matplotlib-inline=0.1.6=py310h06a4308_0
+  - mistune=2.0.4=py310h06a4308_0
+  - nbclassic=1.0.0=py310h06a4308_0
+  - nbclient=0.8.0=py310h06a4308_0
+  - nbconvert=7.10.0=py310h06a4308_0
+  - nbformat=5.9.2=py310h06a4308_0
+  - ncurses=6.4=h6a678d5_0
+  - nest-asyncio=1.5.6=py310h06a4308_0
+  - notebook=6.5.4=py310h06a4308_0
+  - notebook-shim=0.2.3=py310h06a4308_0
+  - openssl=3.0.12=h7f8727e_0
+  - overrides=7.4.0=py310h06a4308_0
+  - pandocfilters=1.5.0=pyhd3eb1b0_0
+  - parso=0.8.3=pyhd3eb1b0_0
+  - pexpect=4.8.0=pyhd3eb1b0_3
+  - pickleshare=0.7.5=pyhd3eb1b0_1003
+  - platformdirs=3.10.0=py310h06a4308_0
+  - prometheus_client=0.14.1=py310h06a4308_0
+  - prompt-toolkit=3.0.36=py310h06a4308_0
+  - ptyprocess=0.7.0=pyhd3eb1b0_2
+  - pure_eval=0.2.2=pyhd3eb1b0_0
+  - pycparser=2.21=pyhd3eb1b0_0
+  - pyopenssl=23.2.0=py310h06a4308_0
+  - pysocks=1.7.1=py310h06a4308_0
+  - python=3.10.13=h955ad1f_0
+  - python-dateutil=2.8.2=pyhd3eb1b0_0
+  - python-fastjsonschema=2.16.2=py310h06a4308_0
+  - python-json-logger=2.0.7=py310h06a4308_0
+  - pytz=2023.3.post1=py310h06a4308_0
+  - pyyaml=6.0.1=py310h5eee18b_0
+  - pyzmq=25.1.0=py310h6a678d5_0
+  - readline=8.2=h5eee18b_0
+  - referencing=0.30.2=py310h06a4308_0
+  - requests=2.31.0=py310h06a4308_0
+  - rfc3339-validator=0.1.4=py310h06a4308_0
+  - rfc3986-validator=0.1.1=py310h06a4308_0
+  - rpds-py=0.10.6=py310hb02cf49_0
+  - send2trash=1.8.2=py310h06a4308_0
+  - setuptools=68.0.0=py310h06a4308_0
+  - six=1.16.0=pyhd3eb1b0_1
+  - soupsieve=2.5=py310h06a4308_0
+  - sqlite=3.41.2=h5eee18b_0
+  - stack_data=0.2.0=pyhd3eb1b0_0
+  - terminado=0.17.1=py310h06a4308_0
+  - tinycss2=1.2.1=py310h06a4308_0
+  - tk=8.6.12=h1ccaba5_0
+  - tomli=2.0.1=py310h06a4308_0
+  - tornado=6.3.3=py310h5eee18b_0
+  - webencodings=0.5.1=py310h06a4308_1
+  - wheel=0.41.2=py310h06a4308_0
+  - y-py=0.5.9=py310h52d8a92_0
+  - yaml=0.2.5=h7b6447c_0
+  - ypy-websocket=0.8.2=py310h06a4308_0
+  - zeromq=4.3.4=h2531618_0
+  - zlib=1.2.13=h5eee18b_0
+  - pip:
+      - absl-py==2.0.0
+      - aiobotocore==2.8.0
+      - aiohttp==3.9.1
+      - aioitertools==0.11.0
+      - aiosignal==1.3.1
+      - altair==5.2.0
+      - annotated-types==0.6.0
+      - antlr4-python3-runtime==4.9.3
+      - anyio==3.7.1
+      - arrow==1.3.0
+      - async-timeout==4.0.3
+      - backoff==2.2.1
+      - bitsandbytes==0.41.3
+      - blessed==1.20.0
+      - boto3==1.33.1
+      - botocore==1.33.1
+      - cachetools==5.3.2
+      - chardet==5.2.0
+      - charset-normalizer==3.3.2
+      - click==8.1.7
+      - colorama==0.4.6
+      - contourpy==1.2.0
+      - croniter==1.4.1
+      - cycler==0.12.1
+      - dateutils==0.6.12
+      - deepdiff==6.7.1
+      - docker==6.1.3
+      - docstring-parser==0.15
+      - exceptiongroup==1.2.0
+      - fastapi==0.104.1
+      - ffmpy==0.3.1
+      - filelock==3.13.1
+      - fonttools==4.46.0
+      - frozenlist==1.4.0
+      - fsspec==2023.12.1
+      - google-auth==2.25.1
+      - google-auth-oauthlib==1.1.0
+      - gradio==4.12.0
+      - gradio-client==0.8.0
+      - grpcio==1.59.3
+      - h11==0.14.0
+      - httpcore==1.0.2
+      - httpx==0.26.0
+      - huggingface-hub==0.19.4
+      - hydra-core==1.3.2
+      - idna==3.6
+      - importlib-resources==6.1.1
+      - inquirer==3.1.4
+      - itsdangerous==2.1.2
+      - jmespath==1.0.1
+      - jsonargparse==4.27.1
+      - kiwisolver==1.4.5
+      - lightning==2.1.2
+      - lightning-api-access==0.0.5
+      - lightning-cloud==0.5.52
+      - lightning-fabric==2.1.2
+      - lightning-utilities==0.10.0
+      - line-profiler==4.1.2
+      - markdown==3.5.1
+      - markdown-it-py==3.0.0
+      - markupsafe==2.1.3
+      - matplotlib==3.8.2
+      - mdurl==0.1.2
+      - mpmath==1.3.0
+      - multidict==6.0.4
+      - networkx==3.2.1
+      - numpy==1.26.2
+      - nvidia-cublas-cu12==12.1.3.1
+      - nvidia-cuda-cupti-cu12==12.1.105
+      - nvidia-cuda-nvrtc-cu12==12.1.105
+      - nvidia-cuda-runtime-cu12==12.1.105
+      - nvidia-cudnn-cu12==8.9.2.26
+      - nvidia-cufft-cu12==11.0.2.54
+      - nvidia-curand-cu12==10.3.2.106
+      - nvidia-cusolver-cu12==11.4.5.107
+      - nvidia-cusparse-cu12==12.1.0.106
+      - nvidia-nccl-cu12==2.18.1
+      - nvidia-nvjitlink-cu12==12.3.101
+      - nvidia-nvtx-cu12==12.1.105
+      - oauthlib==3.2.2
+      - omegaconf==2.3.0
+      - ordered-set==4.1.0
+      - orjson==3.9.10
+      - packaging==23.2
+      - pandas==2.1.3
+      - pillow==10.1.0
+      - protobuf==4.23.4
+      - psutil==5.9.6
+      - pyasn1==0.5.1
+      - pyasn1-modules==0.3.0
+      - pydantic==2.5.2
+      - pydantic-core==2.14.5
+      - pydub==0.25.1
+      - pygments==2.17.2
+      - pyjwt==2.8.0
+      - pyparsing==3.1.1
+      - python-editor==1.0.4
+      - python-multipart==0.0.6
+      - pytorch-lightning==2.1.2
+      - readchar==4.0.5
+      - redis==5.0.1
+      - regex==2023.10.3
+      - requests-oauthlib==1.3.1
+      - rich==13.7.0
+      - rsa==4.9
+      - s3fs==2023.12.1
+      - s3transfer==0.8.0
+      - safetensors==0.4.1
+      - semantic-version==2.10.0
+      - shellingham==1.5.4
+      - sniffio==1.3.0
+      - starlette==0.27.0
+      - starsessions==1.3.0
+      - sympy==1.12
+      - tensorboard==2.15.1
+      - tensorboard-data-server==0.7.2
+      - tensorboardx==2.6.2.2
+      - tokenizers==0.15.0
+      - tomlkit==0.12.0
+      - toolz==0.12.0
+      - torch==2.1.1
+      - torchmetrics==1.2.1
+      - torchvision==0.16.1
+      - tqdm==4.66.1
+      - traitlets==5.14.0
+      - transformers==4.35.2
+      - triton==2.1.0
+      - typer==0.9.0
+      - types-python-dateutil==2.8.19.14
+      - typeshed-client==2.4.0
+      - typing-extensions==4.9.0
+      - tzdata==2023.3
+      - urllib3==2.0.7
+      - uvicorn==0.24.0.post1
+      - wcwidth==0.2.12
+      - websocket-client==1.7.0
+      - websockets==11.0.3
+      - werkzeug==3.0.1
+      - wrapt==1.16.0
+      - yarl==1.9.4

labelmap.py

@@ -1,3 +1,5 @@
+""" Mapping of class IDs to lables. """
+
 DR_LABELMAP = {
     0: 'No DR',
     1: 'Mild',

requirements.txt

@@ -0,0 +1,74 @@
+--extra-index-url https://download.pytorch.org/whl/cpu
+aiofiles==23.2.1
+altair==5.2.0
+annotated-types==0.6.0
+anyio==4.2.0
+attrs==23.2.0
+certifi==2023.11.17
+charset-normalizer==3.3.2
+click==8.1.7
+colorama==0.4.6
+contourpy==1.2.0
+cycler==0.12.1
+exceptiongroup==1.2.0
+fastapi==0.108.0
+ffmpy==0.3.1
+filelock==3.13.1
+fonttools==4.47.0
+fsspec==2023.12.2
+gradio==4.13.0
+gradio_client==0.8.0
+h11==0.14.0
+httpcore==1.0.2
+httpx==0.26.0
+huggingface-hub==0.20.2
+idna==3.6
+importlib-resources==6.1.1
+Jinja2==3.1.2
+jsonschema==4.20.0
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.5
+markdown-it-py==3.0.0
+MarkupSafe==2.1.3
+matplotlib==3.8.2
+mdurl==0.1.2
+mpmath==1.3.0
+networkx==3.2.1
+numpy==1.26.3
+orjson==3.9.10
+packaging==23.2
+pandas==2.1.4
+pillow==10.2.0
+pydantic==2.5.3
+pydantic_core==2.14.6
+pydub==0.25.1
+Pygments==2.17.2
+pyparsing==3.1.1
+python-dateutil==2.8.2
+python-multipart==0.0.6
+pytz==2023.3.post1
+PyYAML==6.0.1
+referencing==0.32.1
+regex==2023.12.25
+requests==2.31.0
+rich==13.7.0
+rpds-py==0.16.2
+safetensors==0.4.1
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.0
+starlette==0.32.0.post1
+sympy==1.12
+tokenizers==0.15.0
+tomlkit==0.12.0
+toolz==0.12.0
+torch==2.1.2+cpu
+tqdm==4.66.1
+transformers==4.36.2
+typer==0.9.0
+typing_extensions==4.9.0
+tzdata==2023.4
+urllib3==2.1.0
+uvicorn==0.25.0
+websockets==11.0.3

train.py

@@ -49,7 +49,15 @@ DataRecord = Tuple[Image.Image, int]
 
 
 class RetinopathyDataset(data.Dataset[DataRecord]):
+    """ A class to access the pre-downloaded Diabetic Retinopathy dataset. """
+
     def __init__(self, data_path: str) -> None:
+        """ Constructor.
+
+        Args:
+            data_path (str): path to the dataset, ex: "retinopathy_data"
+                containing "trainLabels.csv" and "train/".
+        """
         super().__init__()
 
         self.data_path = data_path
@@ -88,21 +96,25 @@ class RetinopathyDataset(data.Dataset[DataRecord]):
         return img_path
 
 
+""" Purpose of a split: training or validation. """
 class Purpose(Enum):
     Train = 0
     Val = 1
 
-
+""" Augmentation transformations for an image and a label. """
 FeatureAndTargetTransforms = Tuple[Callable[..., torch.Tensor],
                                    Callable[..., torch.Tensor]]
 
+""" Feature (image) and target (label) tensors. """
 TensorRecord = Tuple[torch.Tensor, torch.Tensor]
 
-def normalize(arr: np.ndarray) -> np.ndarray:
-    return arr / np.sum(arr)
-
 
 class Split(data.Dataset[TensorRecord], collections.abc.Sequence[TensorRecord]):
+    """ Split is a class that keep a view on a part of a dataset.
+    Split is used to hold the imormation about which samples go to training
+    and which to validation without a need to put these groups of files into
+    separate folders.
+    """
     def __init__(self, dataset: RetinopathyDataset,
                  indices: np.ndarray,
                  purpose: Purpose,
@@ -111,7 +123,24 @@ class Split(data.Dataset[TensorRecord], collections.abc.Sequence[TensorRecord]):
                  stratify_classes: bool = False,
                  use_log_frequencies: bool = False,
                  ):
-
+        """ Constructor.
+
+        Args:
+            dataset (RetinopathyDataset): The dataset on which the Split "views".
+            indices (np.ndarray): Externally provided indices of samples that
+                are "viewed" on.
+            purpose (Purpose): Either train or val, to be able to replicate
+                the data for train split for effecient workers utilization.
+            transforms (FeatureAndTargetTransforms): Functors of feature and
+                target transforms.
+            oversample_factor (int, optional): Expand the training dataset by
+                replication to avoid dataloader stalls on epoch ends. Defaults to 1.
+            stratify_classes (bool, optional): Whether to apply stratified sampling.
+                Defaults to False.
+            use_log_frequencies (bool, optional): If stratify_classes=True,
+                whether to use logarithmic sampling strategy. If False, apply
+                regular even sampling. Defaults to False.
+        """
         self.dataset = dataset
         self.indices = indices
         self.purpose = purpose
@@ -124,22 +153,26 @@ class Split(data.Dataset[TensorRecord], collections.abc.Sequence[TensorRecord]):
         self.per_class_indices: Optional[Dict[int, np.ndarray]] = None
         self.frequencies: Optional[Dict[int, float]] = None
         if self.stratify_classes:
-            self.bucketize_indices()
+            self._bucketize_indices()
             if self.use_log_frequencies:
-                self.calc_frequencies()
+                self._calc_frequencies()
 
-    def calc_frequencies(self):
+    def _calc_frequencies(self):
         assert self.per_class_indices is not None
         counts_dict = {lbl: len(arr) for lbl, arr in self.per_class_indices.items()}
         counts = np.array(list(counts_dict.values()))
-        counts_nrm = normalize(counts)
+        counts_nrm = self._normalize(counts)
         temperature = 50.0 # > 1 to even-out frequencies
-        freqs = normalize(np.log1p(counts_nrm * temperature))
+        freqs = self._normalize(np.log1p(counts_nrm * temperature))
         self.frequencies = {k: freq.item() for k, freq
                             in zip(self.per_class_indices.keys(), freqs)}
         print(self.frequencies)
 
-    def bucketize_indices(self):
+    @staticmethod
+    def _normalize(arr: np.ndarray) -> np.ndarray:
+        return arr / np.sum(arr)
+
+    def _bucketize_indices(self):
         buckets = defaultdict(list)
         for index in self.indices:
             label = self.dataset.get_label_at(index)
@@ -191,6 +224,14 @@ class Split(data.Dataset[TensorRecord], collections.abc.Sequence[TensorRecord]):
                     seed: int = 54,
                     ) -> Tuple['Split', 'Split']:
 
+        """ Prepare train and val splits deterministically.
+
+        Returns:
+            Tuple[Split, Split]:
+                - Train split
+                - Val split
+        """
+
         prng = RandomState(seed)
 
         num_train = int(len(all_data) * train_fraction)
@@ -204,7 +245,8 @@ class Split(data.Dataset[TensorRecord], collections.abc.Sequence[TensorRecord]):
         return train_data, val_data
 
 
-def print_data_stats(dataset: Union[Iterable[DataRecord], DataLoader], split_name: str) -> None:
+def print_data_stats(dataset: Union[Iterable[DataRecord], DataLoader],
+                     split_name: str) -> None:
     labels = []
     for _, label in dataset:
         if isinstance(label, torch.Tensor):
@@ -261,7 +303,16 @@ class Metrics:
         return self
 
 
-def worker_init_fn(worker_id):
+def worker_init_fn(worker_id: int) -> None:
+    """ Initialize workers in a way that they draw different
+    random samples and do not repeat identical pseudorandom
+    sequences of each other, which may be the case with Fork
+    multiprocessing.
+
+    Args:
+        worker_id (int): id of a preprocessing worker process launched
+        by one DDP training process. 
+    """
     state = np.random.get_state()
     assert isinstance(state, tuple)
     assert isinstance(state[1], np.ndarray)
@@ -274,6 +325,7 @@ def worker_init_fn(worker_id):
 
 
 class ViTLightningModule(L.LightningModule):
+    """ Lightning Module that implements neural network training hooks. """
     def __init__(self, debug: bool) -> None:
         super().__init__()
 
@@ -443,6 +495,7 @@ class ViTLightningModule(L.LightningModule):
         return loss
 
     def _dump_train_images(self) -> None:
+        """ Save augmented images to disk for inspection. """
         img_batch, label_batch = next(iter(self._train_dataloader))
         for i_img, (img, label) in enumerate(zip(img_batch, label_batch)):
             img_np = img.cpu().numpy()
@@ -494,18 +547,19 @@ class ViTLightningModule(L.LightningModule):
 
 
 def main():
+    """ Neural network trainer entry point. """
 
     parser = ArgumentParser(description='KAUST-SDAIA Diabetic Retinopathy')
     parser.add_argument('--tag', action='store', type=str,
                         help='Extra suffix to put on the artefact dir name')
-    parser.add_argument('--debug', action='store_true')
+    parser.add_argument('--debug', action='store_true',
+                        help="Dummy training cycle for testing purposes")
     parser.add_argument('--convert-checkpoint', action='store', type=str,
                         help='Convert a checkpoint from training to pickle-independent '
                              'predictor-compatible directory')
 
     args = parser.parse_args()
 
-
     torch.set_float32_matmul_precision('high') # for V100/A100
 
     if args.convert_checkpoint is not None: