Lung Segmentation Using the Montgomery Dataset¶

[Paper] [Notebook] [TF Implementation] [Torch Implementation]

In [1]:

import os
import tempfile
from typing import Any, Dict, List

import cv2
import numpy as np
import pandas as pd
import torch

import fastestimator as fe
from fastestimator.architecture.pytorch import UNet
from fastestimator.dataset.data import montgomery
from fastestimator.op.numpyop import Delete, NumpyOp
from fastestimator.op.numpyop.meta import Sometimes
from fastestimator.op.numpyop.multivariate import HorizontalFlip, Resize, Rotate
from fastestimator.op.numpyop.univariate import Minmax, ReadImage, Reshape
from fastestimator.op.tensorop.loss import CrossEntropy
from fastestimator.op.tensorop.model import ModelOp, UpdateOp
from fastestimator.trace.io import BestModelSaver
from fastestimator.trace.metric import Dice
from fastestimator.util import ImageDisplay, GridDisplay

import os import tempfile from typing import Any, Dict, List import cv2 import numpy as np import pandas as pd import torch import fastestimator as fe from fastestimator.architecture.pytorch import UNet from fastestimator.dataset.data import montgomery from fastestimator.op.numpyop import Delete, NumpyOp from fastestimator.op.numpyop.meta import Sometimes from fastestimator.op.numpyop.multivariate import HorizontalFlip, Resize, Rotate from fastestimator.op.numpyop.univariate import Minmax, ReadImage, Reshape from fastestimator.op.tensorop.loss import CrossEntropy from fastestimator.op.tensorop.model import ModelOp, UpdateOp from fastestimator.trace.io import BestModelSaver from fastestimator.trace.metric import Dice from fastestimator.util import ImageDisplay, GridDisplay

In [2]:

pd.set_option('display.max_colwidth', 500)

pd.set_option('display.max_colwidth', 500)

In [3]:

parameters

batch_size = 4
epochs = 25
train_steps_per_epoch = None
eval_steps_per_epoch = None
save_dir = tempfile.mkdtemp()
data_dir = None

batch_size = 4 epochs = 25 train_steps_per_epoch = None eval_steps_per_epoch = None save_dir = tempfile.mkdtemp() data_dir = None

Download Data¶

We download the Montgomery data first:

In [4]:

csv = montgomery.load_data(root_dir=data_dir)

csv = montgomery.load_data(root_dir=data_dir)

This creates a CSVDataset. Let's see what is inside:

In [5]:

df = pd.DataFrame.from_dict(csv.data, orient='index')

df = pd.DataFrame.from_dict(csv.data, orient='index')

In [6]:

df.head()

df.head()

Out[6]:

	image	mask_left	mask_right
0	MontgomerySet/CXR_png/MCUCXR_0383_1.png	MontgomerySet/ManualMask/leftMask/MCUCXR_0383_1.png	MontgomerySet/ManualMask/rightMask/MCUCXR_0383_1.png
1	MontgomerySet/CXR_png/MCUCXR_0255_1.png	MontgomerySet/ManualMask/leftMask/MCUCXR_0255_1.png	MontgomerySet/ManualMask/rightMask/MCUCXR_0255_1.png
2	MontgomerySet/CXR_png/MCUCXR_0016_0.png	MontgomerySet/ManualMask/leftMask/MCUCXR_0016_0.png	MontgomerySet/ManualMask/rightMask/MCUCXR_0016_0.png
3	MontgomerySet/CXR_png/MCUCXR_0182_1.png	MontgomerySet/ManualMask/leftMask/MCUCXR_0182_1.png	MontgomerySet/ManualMask/rightMask/MCUCXR_0182_1.png
4	MontgomerySet/CXR_png/MCUCXR_0338_1.png	MontgomerySet/ManualMask/leftMask/MCUCXR_0338_1.png	MontgomerySet/ManualMask/rightMask/MCUCXR_0338_1.png

Building Components¶

Now let's set the stage for training:

Step 1: Create Pipeline¶

In [7]:

class CombineLeftRightMask(NumpyOp):
    def forward(self, data: List[np.ndarray], state: Dict[str, Any]) -> List[np.ndarray]:
        mask_left, mask_right = data
        data = np.maximum(mask_left, mask_right)
        return data

class CombineLeftRightMask(NumpyOp): def forward(self, data: List[np.ndarray], state: Dict[str, Any]) -> List[np.ndarray]: mask_left, mask_right = data data = np.maximum(mask_left, mask_right) return data

In [8]:

pipeline = fe.Pipeline(
    train_data=csv,
    eval_data=csv.split(0.2),
    batch_size=batch_size,
    ops=[
        ReadImage(inputs="image", parent_path=csv.parent_path, outputs="image", color_flag="gray"),
        ReadImage(inputs="mask_left", parent_path=csv.parent_path, outputs="mask_left", color_flag="gray", mode='!infer'),
        ReadImage(inputs="mask_right",
                  parent_path=csv.parent_path,
                  outputs="mask_right",
                  color_flag="gray",
                  mode='!infer'),
        CombineLeftRightMask(inputs=("mask_left", "mask_right"), outputs="mask", mode='!infer'),
        Delete(keys=["mask_left", "mask_right"], mode='!infer'),
        Resize(image_in="image", width=512, height=512),
        Resize(image_in="mask", width=512, height=512, mode='!infer'),
        Sometimes(numpy_op=HorizontalFlip(image_in="image", mask_in="mask", mode='train')),
        Sometimes(numpy_op=Rotate(
            image_in="image", mask_in="mask", limit=(-10, 10), border_mode=cv2.BORDER_CONSTANT, mode='train')),
        Minmax(inputs="image", outputs="image"),
        Minmax(inputs="mask", outputs="mask", mode='!infer'),
        Reshape(shape=(1, 512, 512), inputs="image", outputs="image"),
        Reshape(shape=(1, 512, 512), inputs="mask", outputs="mask", mode='!infer')
    ])

pipeline = fe.Pipeline( train_data=csv, eval_data=csv.split(0.2), batch_size=batch_size, ops=[ ReadImage(inputs="image", parent_path=csv.parent_path, outputs="image", color_flag="gray"), ReadImage(inputs="mask_left", parent_path=csv.parent_path, outputs="mask_left", color_flag="gray", mode='!infer'), ReadImage(inputs="mask_right", parent_path=csv.parent_path, outputs="mask_right", color_flag="gray", mode='!infer'), CombineLeftRightMask(inputs=("mask_left", "mask_right"), outputs="mask", mode='!infer'), Delete(keys=["mask_left", "mask_right"], mode='!infer'), Resize(image_in="image", width=512, height=512), Resize(image_in="mask", width=512, height=512, mode='!infer'), Sometimes(numpy_op=HorizontalFlip(image_in="image", mask_in="mask", mode='train')), Sometimes(numpy_op=Rotate( image_in="image", mask_in="mask", limit=(-10, 10), border_mode=cv2.BORDER_CONSTANT, mode='train')), Minmax(inputs="image", outputs="image"), Minmax(inputs="mask", outputs="mask", mode='!infer'), Reshape(shape=(1, 512, 512), inputs="image", outputs="image"), Reshape(shape=(1, 512, 512), inputs="mask", outputs="mask", mode='!infer') ])

Let's see if the Pipeline output is reasonable. We call get_results to get outputs from Pipeline.

In [9]:

batch_data = pipeline.get_results()

batch_data = pipeline.get_results()

In [10]:

batch_index = 1
GridDisplay([ImageDisplay(image=np.squeeze(batch_data['image'][batch_index])), 
             ImageDisplay(image=np.squeeze(batch_data['mask'][batch_index]))
            ]).show()

batch_index = 1 GridDisplay([ImageDisplay(image=np.squeeze(batch_data['image'][batch_index])), ImageDisplay(image=np.squeeze(batch_data['mask'][batch_index])) ]).show()

Step 2: Create Network¶

In [11]:

model = fe.build(
    model_fn=lambda: UNet(input_size=(1, 512, 512)),
    optimizer_fn=lambda x: torch.optim.Adam(params=x, lr=0.0001),
    model_name="lung_segmentation"
)

model = fe.build( model_fn=lambda: UNet(input_size=(1, 512, 512)), optimizer_fn=lambda x: torch.optim.Adam(params=x, lr=0.0001), model_name="lung_segmentation" )

In [12]:

network = fe.Network(ops=[
    ModelOp(inputs="image", model=model, outputs="pred_segment"),
    CrossEntropy(inputs=("pred_segment", "mask"), outputs="loss", form="binary"),
    UpdateOp(model=model, loss_name="loss")
])

network = fe.Network(ops=[ ModelOp(inputs="image", model=model, outputs="pred_segment"), CrossEntropy(inputs=("pred_segment", "mask"), outputs="loss", form="binary"), UpdateOp(model=model, loss_name="loss") ])

Step 3: Create Estimator¶

In [13]:

traces = [
    Dice(true_key="mask", pred_key="pred_segment"),
    BestModelSaver(model=model, save_dir=save_dir, metric='Dice', save_best_mode='max')
]

traces = [ Dice(true_key="mask", pred_key="pred_segment"), BestModelSaver(model=model, save_dir=save_dir, metric='Dice', save_best_mode='max') ]

In [14]:

estimator = fe.Estimator(network=network,
                         pipeline=pipeline,
                         epochs=epochs,
                         log_steps=20,
                         traces=traces,
                         train_steps_per_epoch=train_steps_per_epoch,
                         eval_steps_per_epoch=eval_steps_per_epoch)

estimator = fe.Estimator(network=network, pipeline=pipeline, epochs=epochs, log_steps=20, traces=traces, train_steps_per_epoch=train_steps_per_epoch, eval_steps_per_epoch=eval_steps_per_epoch)

Training¶

In [15]:

estimator.fit()

estimator.fit()

    ______           __  ______     __  _                 __            
   / ____/___ ______/ /_/ ____/____/ /_(_)___ ___  ____ _/ /_____  _____
  / /_  / __ `/ ___/ __/ __/ / ___/ __/ / __ `__ \/ __ `/ __/ __ \/ ___/
 / __/ / /_/ (__  ) /_/ /___(__  ) /_/ / / / / / / /_/ / /_/ /_/ / /    
/_/    \__,_/____/\__/_____/____/\__/_/_/ /_/ /_/\__,_/\__/\____/_/     
                                                                        

FastEstimator-Start: step: 1; logging_interval: 20; num_device: 1;
FastEstimator-Train: step: 1; loss: 0.61223024;
FastEstimator-Train: step: 20; loss: 0.3766136; steps/sec: 7.9;
FastEstimator-Train: step: 28; epoch: 1; epoch_time: 11.86 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 7.82;
Eval Progress: 4/7; steps/sec: 13.23;
Eval Progress: 7/7; steps/sec: 16.75;
FastEstimator-BestModelSaver: Saved model to /tmp/tmp22fq9j1i/lung_segmentation_best_Dice.pt
FastEstimator-Eval: step: 28; epoch: 1; Dice: 0.8265303; loss: 0.2719451; max_Dice: 0.8265303; since_best_Dice: 0;
FastEstimator-Train: step: 40; loss: 0.14792308; steps/sec: 1.8;
FastEstimator-Train: step: 56; epoch: 2; epoch_time: 11.99 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 10.29;
Eval Progress: 4/7; steps/sec: 13.7;
Eval Progress: 7/7; steps/sec: 10.47;
FastEstimator-BestModelSaver: Saved model to /tmp/tmp22fq9j1i/lung_segmentation_best_Dice.pt
FastEstimator-Eval: step: 56; epoch: 2; Dice: 0.9142313; loss: 0.12388889; max_Dice: 0.9142313; since_best_Dice: 0;
FastEstimator-Train: step: 60; loss: 0.22085804; steps/sec: 1.87;
FastEstimator-Train: step: 80; loss: 0.13877149; steps/sec: 8.15;
FastEstimator-Train: step: 84; epoch: 3; epoch_time: 11.61 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 4.56;
Eval Progress: 4/7; steps/sec: 11.72;
Eval Progress: 7/7; steps/sec: 18.49;
FastEstimator-Eval: step: 84; epoch: 3; Dice: 0.8878712; loss: 0.13381886; max_Dice: 0.9142313; since_best_Dice: 1;
FastEstimator-Train: step: 100; loss: 0.09233445; steps/sec: 1.8;
FastEstimator-Train: step: 112; epoch: 4; epoch_time: 12.2 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 8.0;
Eval Progress: 4/7; steps/sec: 11.08;
Eval Progress: 7/7; steps/sec: 10.61;
FastEstimator-BestModelSaver: Saved model to /tmp/tmp22fq9j1i/lung_segmentation_best_Dice.pt
FastEstimator-Eval: step: 112; epoch: 4; Dice: 0.93393403; loss: 0.09484797; max_Dice: 0.93393403; since_best_Dice: 0;
FastEstimator-Train: step: 120; loss: 0.08613084; steps/sec: 1.83;
FastEstimator-Train: step: 140; loss: 0.06546959; steps/sec: 5.96;
FastEstimator-Train: step: 140; epoch: 5; epoch_time: 11.91 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 7.08;
Eval Progress: 4/7; steps/sec: 16.35;
Eval Progress: 7/7; steps/sec: 9.35;
FastEstimator-BestModelSaver: Saved model to /tmp/tmp22fq9j1i/lung_segmentation_best_Dice.pt
FastEstimator-Eval: step: 140; epoch: 5; Dice: 0.9473874; loss: 0.069016606; max_Dice: 0.9473874; since_best_Dice: 0;
FastEstimator-Train: step: 160; loss: 0.050840445; steps/sec: 1.97;
FastEstimator-Train: step: 168; epoch: 6; epoch_time: 11.94 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 1.92;
Eval Progress: 4/7; steps/sec: 10.57;
Eval Progress: 7/7; steps/sec: 13.94;
FastEstimator-BestModelSaver: Saved model to /tmp/tmp22fq9j1i/lung_segmentation_best_Dice.pt
FastEstimator-Eval: step: 168; epoch: 6; Dice: 0.9499039; loss: 0.066906914; max_Dice: 0.9499039; since_best_Dice: 0;
FastEstimator-Train: step: 180; loss: 0.05280122; steps/sec: 1.83;
FastEstimator-Train: step: 196; epoch: 7; epoch_time: 11.97 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 10.8;
Eval Progress: 4/7; steps/sec: 14.02;
Eval Progress: 7/7; steps/sec: 15.59;
FastEstimator-BestModelSaver: Saved model to /tmp/tmp22fq9j1i/lung_segmentation_best_Dice.pt
FastEstimator-Eval: step: 196; epoch: 7; Dice: 0.9537325; loss: 0.064191505; max_Dice: 0.9537325; since_best_Dice: 0;
FastEstimator-Train: step: 200; loss: 0.08146733; steps/sec: 1.78;
FastEstimator-Train: step: 220; loss: 0.06502332; steps/sec: 8.16;
FastEstimator-Train: step: 224; epoch: 8; epoch_time: 12.21 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 8.1;
Eval Progress: 4/7; steps/sec: 12.47;
Eval Progress: 7/7; steps/sec: 10.71;
FastEstimator-BestModelSaver: Saved model to /tmp/tmp22fq9j1i/lung_segmentation_best_Dice.pt
FastEstimator-Eval: step: 224; epoch: 8; Dice: 0.9608661; loss: 0.053247754; max_Dice: 0.9608661; since_best_Dice: 0;
FastEstimator-Train: step: 240; loss: 0.04226496; steps/sec: 1.82;
FastEstimator-Train: step: 252; epoch: 9; epoch_time: 12.04 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 11.38;
Eval Progress: 4/7; steps/sec: 13.69;
Eval Progress: 7/7; steps/sec: 14.5;
FastEstimator-BestModelSaver: Saved model to /tmp/tmp22fq9j1i/lung_segmentation_best_Dice.pt
FastEstimator-Eval: step: 252; epoch: 9; Dice: 0.9623051; loss: 0.049131673; max_Dice: 0.9623051; since_best_Dice: 0;
FastEstimator-Train: step: 260; loss: 0.03119469; steps/sec: 1.91;
FastEstimator-Train: step: 280; loss: 0.04104006; steps/sec: 5.95;
FastEstimator-Train: step: 280; epoch: 10; epoch_time: 11.39 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 7.46;
Eval Progress: 4/7; steps/sec: 13.65;
Eval Progress: 7/7; steps/sec: 11.05;
FastEstimator-Eval: step: 280; epoch: 10; Dice: 0.960547; loss: 0.05416977; max_Dice: 0.9623051; since_best_Dice: 1;
FastEstimator-Train: step: 300; loss: 0.047458448; steps/sec: 1.91;
FastEstimator-Train: step: 308; epoch: 11; epoch_time: 12.33 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 8.09;
Eval Progress: 4/7; steps/sec: 14.3;
Eval Progress: 7/7; steps/sec: 15.39;
FastEstimator-BestModelSaver: Saved model to /tmp/tmp22fq9j1i/lung_segmentation_best_Dice.pt
FastEstimator-Eval: step: 308; epoch: 11; Dice: 0.9653207; loss: 0.04500082; max_Dice: 0.9653207; since_best_Dice: 0;
FastEstimator-Train: step: 320; loss: 0.03846242; steps/sec: 1.84;
FastEstimator-Train: step: 336; epoch: 12; epoch_time: 11.88 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 6.19;
Eval Progress: 4/7; steps/sec: 10.46;
Eval Progress: 7/7; steps/sec: 16.47;
FastEstimator-BestModelSaver: Saved model to /tmp/tmp22fq9j1i/lung_segmentation_best_Dice.pt
FastEstimator-Eval: step: 336; epoch: 12; Dice: 0.9654893; loss: 0.044765513; max_Dice: 0.9654893; since_best_Dice: 0;
FastEstimator-Train: step: 340; loss: 0.037619226; steps/sec: 1.84;
FastEstimator-Train: step: 360; loss: 0.035972364; steps/sec: 8.24;
FastEstimator-Train: step: 364; epoch: 13; epoch_time: 11.82 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 1.54;
Eval Progress: 4/7; steps/sec: 11.95;
Eval Progress: 7/7; steps/sec: 12.2;
FastEstimator-Eval: step: 364; epoch: 13; Dice: 0.96202266; loss: 0.0538786; max_Dice: 0.9654893; since_best_Dice: 1;
FastEstimator-Train: step: 380; loss: 0.03339825; steps/sec: 1.8;
FastEstimator-Train: step: 392; epoch: 14; epoch_time: 11.93 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 1.77;
Eval Progress: 4/7; steps/sec: 10.01;
Eval Progress: 7/7; steps/sec: 13.33;
FastEstimator-BestModelSaver: Saved model to /tmp/tmp22fq9j1i/lung_segmentation_best_Dice.pt
FastEstimator-Eval: step: 392; epoch: 14; Dice: 0.96931; loss: 0.043865643; max_Dice: 0.96931; since_best_Dice: 0;
FastEstimator-Train: step: 400; loss: 0.031866133; steps/sec: 1.79;
FastEstimator-Train: step: 420; loss: 0.07556894; steps/sec: 5.8;
FastEstimator-Train: step: 420; epoch: 15; epoch_time: 12.42 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 2.77;
Eval Progress: 4/7; steps/sec: 14.6;
Eval Progress: 7/7; steps/sec: 14.13;
FastEstimator-Eval: step: 420; epoch: 15; Dice: 0.9669868; loss: 0.045212775; max_Dice: 0.96931; since_best_Dice: 1;
FastEstimator-Train: step: 440; loss: 0.06996163; steps/sec: 1.96;
FastEstimator-Train: step: 448; epoch: 16; epoch_time: 12.0 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 9.26;
Eval Progress: 4/7; steps/sec: 4.09;
Eval Progress: 7/7; steps/sec: 9.05;
FastEstimator-Eval: step: 448; epoch: 16; Dice: 0.96843016; loss: 0.043515813; max_Dice: 0.96931; since_best_Dice: 2;
FastEstimator-Train: step: 460; loss: 0.04542666; steps/sec: 1.77;
FastEstimator-Train: step: 476; epoch: 17; epoch_time: 12.26 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 7.92;
Eval Progress: 4/7; steps/sec: 16.38;
Eval Progress: 7/7; steps/sec: 13.81;
FastEstimator-Eval: step: 476; epoch: 17; Dice: 0.9618596; loss: 0.049392406; max_Dice: 0.96931; since_best_Dice: 3;
FastEstimator-Train: step: 480; loss: 0.10188608; steps/sec: 1.88;
FastEstimator-Train: step: 500; loss: 0.06218755; steps/sec: 8.36;
FastEstimator-Train: step: 504; epoch: 18; epoch_time: 11.62 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 7.89;
Eval Progress: 4/7; steps/sec: 12.78;
Eval Progress: 7/7; steps/sec: 11.01;
FastEstimator-BestModelSaver: Saved model to /tmp/tmp22fq9j1i/lung_segmentation_best_Dice.pt
FastEstimator-Eval: step: 504; epoch: 18; Dice: 0.9720732; loss: 0.03817087; max_Dice: 0.9720732; since_best_Dice: 0;
FastEstimator-Train: step: 520; loss: 0.08881314; steps/sec: 1.86;
FastEstimator-Train: step: 532; epoch: 19; epoch_time: 11.7 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 7.32;
Eval Progress: 4/7; steps/sec: 11.15;
Eval Progress: 7/7; steps/sec: 8.55;
FastEstimator-BestModelSaver: Saved model to /tmp/tmp22fq9j1i/lung_segmentation_best_Dice.pt
FastEstimator-Eval: step: 532; epoch: 19; Dice: 0.97263706; loss: 0.03696952; max_Dice: 0.97263706; since_best_Dice: 0;
FastEstimator-Train: step: 540; loss: 0.02760274; steps/sec: 1.78;
FastEstimator-Train: step: 560; loss: 0.031888127; steps/sec: 5.93;
FastEstimator-Train: step: 560; epoch: 20; epoch_time: 12.29 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 7.17;
Eval Progress: 4/7; steps/sec: 10.18;
Eval Progress: 7/7; steps/sec: 12.52;
FastEstimator-Eval: step: 560; epoch: 20; Dice: 0.9725167; loss: 0.036553062; max_Dice: 0.97263706; since_best_Dice: 1;
FastEstimator-Train: step: 580; loss: 0.02704905; steps/sec: 1.96;
FastEstimator-Train: step: 588; epoch: 21; epoch_time: 12.08 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 7.61;
Eval Progress: 4/7; steps/sec: 12.06;
Eval Progress: 7/7; steps/sec: 16.31;
FastEstimator-Eval: step: 588; epoch: 21; Dice: 0.96258086; loss: 0.05462244; max_Dice: 0.97263706; since_best_Dice: 2;
FastEstimator-Train: step: 600; loss: 0.047244675; steps/sec: 1.86;
FastEstimator-Train: step: 616; epoch: 22; epoch_time: 11.64 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 1.21;
Eval Progress: 4/7; steps/sec: 11.61;
Eval Progress: 7/7; steps/sec: 12.71;
FastEstimator-BestModelSaver: Saved model to /tmp/tmp22fq9j1i/lung_segmentation_best_Dice.pt
FastEstimator-Eval: step: 616; epoch: 22; Dice: 0.9769085; loss: 0.031764; max_Dice: 0.9769085; since_best_Dice: 0;
FastEstimator-Train: step: 620; loss: 0.029085426; steps/sec: 1.81;
FastEstimator-Train: step: 640; loss: 0.03575117; steps/sec: 7.88;
FastEstimator-Train: step: 644; epoch: 23; epoch_time: 12.17 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 7.77;
Eval Progress: 4/7; steps/sec: 14.03;
Eval Progress: 7/7; steps/sec: 13.28;
FastEstimator-Eval: step: 644; epoch: 23; Dice: 0.97067636; loss: 0.039316583; max_Dice: 0.9769085; since_best_Dice: 1;
FastEstimator-Train: step: 660; loss: 0.021667594; steps/sec: 1.84;
FastEstimator-Train: step: 672; epoch: 24; epoch_time: 11.72 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 4.98;
Eval Progress: 4/7; steps/sec: 7.56;
Eval Progress: 7/7; steps/sec: 12.22;
FastEstimator-Eval: step: 672; epoch: 24; Dice: 0.9691401; loss: 0.04364162; max_Dice: 0.9769085; since_best_Dice: 2;
FastEstimator-Train: step: 680; loss: 0.05466278; steps/sec: 1.84;
FastEstimator-Train: step: 700; loss: 0.034504842; steps/sec: 6.04;
FastEstimator-Train: step: 700; epoch: 25; epoch_time: 11.94 sec;
Eval Progress: 1/7;
Eval Progress: 2/7; steps/sec: 14.17;
Eval Progress: 4/7; steps/sec: 17.96;
Eval Progress: 7/7; steps/sec: 17.99;
FastEstimator-Eval: step: 700; epoch: 25; Dice: 0.974748; loss: 0.033991832; max_Dice: 0.9769085; since_best_Dice: 3;
FastEstimator-Finish: step: 700; lung_segmentation_lr: 0.0001; total_time: 528.8 sec;

Inferencing¶

Let's visualize the prediction from the neural network. We select a random image from the dataset:

In [16]:

image_path = df['image'].sample(random_state=3).values[0]

image_path = df['image'].sample(random_state=3).values[0]

Pass the image through Pipeline and Network¶

We create a data dict, and call Pipeline.transform().

In [17]:

data = {'image': image_path}
data = pipeline.transform(data, mode="infer")

data = {'image': image_path} data = pipeline.transform(data, mode="infer")

After the Pipeline, we rebuild our model by providing the trained weights path and pass it to a new Network:

In [18]:

weights_path = os.path.join(save_dir, "lung_segmentation_best_Dice.pt") # your model_path

model = fe.build(model_fn=lambda: UNet(input_size=(1, 512, 512)),
                 optimizer_fn=lambda x: torch.optim.Adam(params=x, lr=0.0001),
                 model_name="lung_segmentation",
                 weights_path=weights_path)

weights_path = os.path.join(save_dir, "lung_segmentation_best_Dice.pt") # your model_path model = fe.build(model_fn=lambda: UNet(input_size=(1, 512, 512)), optimizer_fn=lambda x: torch.optim.Adam(params=x, lr=0.0001), model_name="lung_segmentation", weights_path=weights_path)

In [19]:

network = fe.Network(ops=[ModelOp(inputs="image", model=model, outputs="pred_segment")])

network = fe.Network(ops=[ModelOp(inputs="image", model=model, outputs="pred_segment")])

We call Network.transform() to get outputs from our Network:

In [20]:

pred = network.transform(data, mode="infer")

pred = network.transform(data, mode="infer")

Visualize Outputs¶

In [21]:

img = np.squeeze(pred['image'].numpy())
mask = np.squeeze(1.0 * (pred['pred_segment'] > 0.5))

GridDisplay([ImageDisplay(image=img, title="Original Lung"), 
             ImageDisplay(image=img, masks=mask, title="Prediction Mask")]).show()

img = np.squeeze(pred['image'].numpy()) mask = np.squeeze(1.0 * (pred['pred_segment'] > 0.5)) GridDisplay([ImageDisplay(image=img, title="Original Lung"), ImageDisplay(image=img, masks=mask, title="Prediction Mask")]).show()