Partial docstrings

README.md

@@ -1,8 +1,30 @@
-# diabetic-retinopathy-detection
+# Diabetic Retinopathy Detection with AI
 
-# Installataion
+## Setup
+
+### Gradio app environment
+
+TODO
+
+### Training environment
 
 Create conda environment from YAML:
 ```bash
-mamba env create -n retinopathy_restore -f environment.yml
-```
+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/"
 
@@ -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

@@ -69,7 +69,6 @@ dependencies:
  - parso=0.8.3=pyhd3eb1b0_0
  - pexpect=4.8.0=pyhd3eb1b0_3
  - pickleshare=0.7.5=pyhd3eb1b0_1003
- - pip=23.3.1=py310h06a4308_0
  - platformdirs=3.10.0=py310h06a4308_0
  - prometheus_client=0.14.1=py310h06a4308_0
  - prompt-toolkit=3.0.36=py310h06a4308_0
@@ -104,7 +103,6 @@ dependencies:
  - tornado=6.3.3=py310h5eee18b_0
  - webencodings=0.5.1=py310h06a4308_1
  - wheel=0.41.2=py310h06a4308_0
- - xz=5.4.5=h5eee18b_0
  - y-py=0.5.9=py310h52d8a92_0
  - yaml=0.2.5=h7b6447c_0
  - ypy-websocket=0.8.2=py310h06a4308_0

labelmap.py

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

requirements.txt

@@ -0,0 +1,2 @@
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch==2.1.2+cpu

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: