num_workers = 0 for cpu

f068dd86 · s193396 · b4ea1bf3 · f068dd86
Commit f068dd86 authored 4 months ago by s193396
--- a/docs/notebooks/UNet.ipynb
+++ b/docs/notebooks/UNet.ipynb
@@ -459,7 +459,8 @@
    "train_loader, val_loader, test_loader = qim3d.ml.prepare_dataloaders(train_set, \n",
    "                                                                     val_set,\n",
    "                                                                     test_set,\n",
-    "                                                                     batch_size = 1)"
+    "                                                                     batch_size = 1,\n",
+    "                                                                     num_workers = 0)"
   ]
  },
  {

 %% Cell type:markdown id: tags:
 # **Deep Learning Volume Segmentation (3D UNet)**
 Authors: Anna Ekner (s193396@dtu.dk)
 This notebook aims to demonstrate the feasibility of implementing a comprehensive deep learning segmentation pipeline solely leveraging the capabilities offered by the `qim3d` library. Specifically, it highlights the use of the synthetic data generation functionalities to create a volumetric dataset with associated labels, and walks through the process of creating and training a 3D UNet model using this synthetic dataset.
 %% Cell type:code id: tags:
 ``` python
 import qim3d
 import glob
 import os
 import numpy as np
 ```
 %% Cell type:markdown id: tags:
 ### **1. Generate synthetic dataset**
 %% Cell type:markdown id: tags:
 #### 1.1 Example of data sample (probably should be after creating the dataset??)
 %% Cell type:markdown id: tags:
 Synthetic dataset and associated labels.
 %% Cell type:code id: tags:
 ``` python
 # num_objects = 5
 # vol, labels = qim3d.generate.noise_object_collection(
 #     num_objects = num_objects,
 #     collection_shape = (128, 128, 128),
 #     min_object_noise = 0.03,
 #     max_object_noise = 0.08,
 #     )
 ```
 %% Cell type:code id: tags:
 ``` python
 # Visualize synthetic collection
 # qim3d.viz.volumetric(vol)
 ```
 %% Cell type:code id: tags:
 ``` python
 # Visualize slices
 # qim3d.viz.slicer(vol)
 ```
 %% Cell type:markdown id: tags:
 There will be $N + 1$ unique labels, because one extra for background.
 But we want only 2 labels: foreground and background.
 %% Cell type:code id: tags:
 ``` python
 # # Convert N + 1 labels into 2 labels (background and object)
 # labels = (labels > 0).astype(int)
 ```
 %% Cell type:code id: tags:
 ``` python
 # Visualize labels
 # qim3d.viz.slicer(labels)
 ```
 %% Cell type:markdown id: tags:
 #### 1.2 Create folder structure
 %% Cell type:code id: tags:
 ``` python
 # Base path for the training data
 base_path = os.path.expanduser("~/dataset")
 # Create directories
 print("Creating directories:")
 for folder_split in ["train", "test"]:
    for folder_type in ["images", "labels"]:
        path = os.path.join(base_path, folder_split, folder_type)
        os.makedirs(path, exist_ok=True)
        print(path)
 # Here we have the option to remove any previous files
 clean_files = True
 if clean_files:
    for root, dirs, files in os.walk(base_path):
        for file in files:
            file_path = os.path.join(root, file)
            os.remove(file_path)
 ```
 %% Output
    Creating directories:
    C:\Users\s193396/dataset\train\images
    C:\Users\s193396/dataset\train\labels
    C:\Users\s193396/dataset\test\images
    C:\Users\s193396/dataset\test\labels
 %% Cell type:markdown id: tags:
 #### 1.3 Create dataset
 %% Cell type:markdown id: tags:
 We need to create a dataset of multiple volumes
 %% Cell type:code id: tags:
 ``` python
 # Specify the number of training and testing samples
 num_train = 4
 num_test = 1
 ```
 %% Cell type:code id: tags:
 ``` python
 # Specify the number of objects in each volume
 num_objects = 5
 # Generate training and testing samples
 for idx in range(num_train + num_test):
    # Determine the folder split (train or test)
    if idx < num_train:
        folder_split = "train"
    else:
        folder_split = "test"
    # TODO: Figure out whether or not the seed makes it such that all volumes are identical?
    vol, label = qim3d.generate.noise_object_collection(
        num_objects=num_objects,
        collection_shape=(128, 128, 128),
        min_object_noise=0.03,
        max_object_noise=0.08,
    )
    # Convert N + 1 labels into 2 labels (background and object)
    label = (label > 0).astype(int)
    # Save volume
    qim3d.io.save(os.path.join(base_path, folder_split, "images", f"{idx + 1}.nii.gz"), vol, compression = True, replace=True)
    # Save label
    qim3d.io.save(os.path.join(base_path, folder_split, "labels", f"{idx + 1}.nii.gz"), label, compression = True, replace=True)
 ```
 %% Output
 %% Cell type:code id: tags:
 ``` python
 # Check the image indices for training and testing data
 for folder_split in ["train", "test"]:
    path = os.path.join(base_path, folder_split, "images")
    files = os.listdir(path)
    files = [int(os.path.basename(f)[0]) for f in files]
    if folder_split == "train":
        print(f"Image indices for training data.....: {files}")
    elif folder_split == "test":
        print(f"Image indices for testing data......: {files}")
 ```
 %% Output
    Image indices for training data.....: [1, 2, 3, 4]
    Image indices for testing data......: [5]
 %% Cell type:markdown id: tags:
 ### **2. Build 3D UNet model**
 %% Cell type:markdown id: tags:
 #### 2.1 Instantiate UNet model
 %% Cell type:code id: tags:
 ``` python
 model = qim3d.ml.models.UNet(size = 'small', dropout = 0.25)
 ```
 %% Cell type:markdown id: tags:
 #### 2.2 Define augmentations
 %% Cell type:code id: tags:
 ``` python
 augmentation =  qim3d.ml.Augmentation(resize = 'crop', transform_train = 'light')
 ```
 %% Cell type:markdown id: tags:
 #### 2.3 Divide dataset into train and test splits
 %% Cell type:code id: tags:
 ``` python
 # # Debug stuff
 # import qim3d
 # import numpy as np
 # vol, labels = qim3d.generate.noise_object_collection(
 #     num_objects = 1,
 #     collection_shape = (128, 128, 128),
 #     min_object_noise = 0.03,
 #     max_object_noise = 0.04,
 #     )
 # augmentation =  qim3d.ml.Augmentation(resize = 'crop', transform_train = 'light')
 # print(f"Before augmentation: {vol.shape} and {vol.dtype}")
 # print(f"Augmentation: {augmentation}\n")
 # compose_augmentation = augmentation.augment(*vol.shape, level = augmentation.transform_train)
 # vol_expanded = np.expand_dims(vol, axis = 0)
 # transformed_vol = compose_augmentation(vol_expanded)
 # print(f"After channel expansion and augmentation: {transformed_vol.shape}")
 # print(f"Dtype after augmentation: {transformed_vol.dtype}\n")
 # # Remove the channel dimension
 # vol_reduced = np.squeeze(transformed_vol)
 # print(f"After reducing channel dimension: {vol_reduced.shape}")
 # # qim3d.viz.volumetric(vol_reduced)
 ```
 %% Output
    Before augmentation: (128, 128, 128) and uint8
    Augmentation: <qim3d.ml._augmentations.Augmentation object at 0x0000028B367D1870>
    After channel expansion and augmentation: torch.Size([1, 128, 128, 128])
    Dtype after augmentation: torch.uint8
    After reducing channel dimension: (128, 128, 128)
 %% Cell type:markdown id: tags:
 Stuff
 %% Cell type:code id: tags:
 ``` python
 # datasets and dataloaders
 train_set, val_set, test_set = qim3d.ml.prepare_datasets(path = base_path,
                                                         val_fraction = 0.5,
                                                         model = model,
                                                         augmentation = augmentation)
 train_loader, val_loader, test_loader = qim3d.ml.prepare_dataloaders(train_set,
                                                                     val_set,
                                                                     test_set,
-                                                                     batch_size = 1)
+                                                                     batch_size = 1,
+                                                                     num_workers = 0)
 ```
 %% Output
    3D image shape:  (128, 128, 128)
 %% Cell type:markdown id: tags:
 ### **3. Train model**
 %% Cell type:markdown id: tags:
 #### 3.1 Define training hyperparameters
 %% Cell type:code id: tags:
 ``` python
 # hyperparameters
 hyperparameters = qim3d.ml.Hyperparameters(model, n_epochs=10,
                                           learning_rate = 5e-3, loss_function='DiceCE',
                                           weight_decay=1e-3)
 ```
 %% Cell type:markdown id: tags:
 #### 3.2 Train model
 %% Cell type:code id: tags:
 ``` python
 # training model
 qim3d.ml.train_model(model, hyperparameters, train_loader, val_loader, plot=True)
 ```
 %% Output
    Epoch   0, train loss: 1.6603, val loss: 1.3486
    Epoch   5, train loss: 1.4026, val loss: 0.4789
 %% Cell type:markdown id: tags:
 ### **4. Test model**
 %% Cell type:code id: tags:
 ``` python
 # Needs to be updated to handle 3D as well
 in_targ_preds_test = qim3d.ml.inference(test_set, model)
 ```
 %% Output
    ---------------------------------------------------------------------------
    ValueError                                Traceback (most recent call last)
 Cell     In[24], line 1
    ----> 1 in_targ_preds_test = qim3d.ml.inference(test_set, model)
 File     c:\s193396\qim3d\qim3d\ml\_ml_utils.py:214, in inference(data, model)
        212     pass
        213 else:
    --> 214     raise ValueError("Input image must be (C,H,W) format")
        216 model.to(device)
        217 model.eval()
    ValueError: Input image must be (C,H,W) format