monai_crf_vs_simplecrf.py

# Install required packages:
# !pip -q install simplecrf nibabel wget
# !BUILD_MONAI=1 pip -q install git+https://github.com/Project-MONAI/MONAI#egg=monai

from monai.networks.blocks import CRF

import wget
import nibabel
import denseCRF3D
import numpy as np
import torch

from functools import wraps
from time import time

# define a simple timing decorator, use for both simplecrf's densecrf3d and monai's CRF
def timing(f):
    @wraps(f)
    def wrap(*args, **kw):
        ts = time()
        result = f(*args, **kw)
        te = time()
        print('func:%r took: %2.4f sec' % (f.__name__, te-ts))
        return result
    return wrap
    
# run functions to get timings on the two functions
# densecrf3d run with timing
@timing
def densecrf3d(I, P, param):
    return denseCRF3D.densecrf3d(I, P, param)

# MONAI's CRF run with timing
@timing
def runmonaicrf(I, P, crf_func):
  return crf_func(P, I)

def load_sample_data():
    # download sample data used in this example from github simplecrf repository
    wget.download('https://github.com/HiLab-git/SimpleCRF/raw/master/data/2013_12_1_img.nii.gz')
    wget.download('https://github.com/HiLab-git/SimpleCRF/raw/master/data/2013_12_1_init.nii.gz')

    # Load some sample data
    I1Nii = nibabel.load('2013_12_1_img.nii.gz')
    PNii  = nibabel.load('2013_12_1_init.nii.gz')
    I1 = I1Nii.get_data()
    P  = PNii.get_data()

    # convert input to intenstiy range of [0, 255]
    I = np.asarray([I1], np.float32)
    I = np.transpose(I, [1, 2, 3, 0])
    I = I / I.max()* 255
    I = np.asarray(I, np.uint8)

    # probability map for each class
    P = 0.5 + (P - 0.5) * 0.8
    P = np.asarray([1.0 - P, P], np.float32)
    P = np.transpose(P, [1, 2, 3, 0])

    return I, P

if __name__ == '__main__':
    I, P = load_sample_data()

    print(I.shape)
    print(P.shape)

    # define same number of iterations to run for each method
    ITERATIONS = 20
    
    # run both methods on cpu

    # 1. SimpleCRF's CRF with ITERATIONS
    dense_crf_param = {}
    dense_crf_param['MaxIterations'] = ITERATIONS
    dense_crf_param['PosW'] = 3
    dense_crf_param['PosRStd'] = 1
    dense_crf_param['PosCStd'] = 1
    dense_crf_param['PosZStd'] = 1
    dense_crf_param['BilateralW'] = 5
    dense_crf_param['BilateralRStd'] = 5
    dense_crf_param['BilateralCStd'] = 5
    dense_crf_param['BilateralZStd'] = 5
    dense_crf_param['ModalityNum'] = 1
    dense_crf_param['BilateralModsStds'] = (5,)

    lab = densecrf3d(I, P, dense_crf_param)

    # 2. MONAI's CRF with ITERATIONS - CPU
    device='cpu'
    I_pt = np.moveaxis(I, source=-1, destination=0)
    P_pt = np.moveaxis(P, source=-1, destination=0)
    I_pt = np.expand_dims(I_pt, axis=0)
    P_pt = np.expand_dims(P_pt, axis=0)
    I_pt = torch.from_numpy(I_pt).to(dtype=torch.float, device=torch.device(device))
    P_pt = torch.from_numpy(P_pt).to(dtype=torch.float, device=torch.device(device))

    crf_func = CRF(
        bilateral_weight=5.0,
        gaussian_weight=3.0,
        bilateral_spatial_sigma=1.0,
        bilateral_color_sigma=5.0,
        gaussian_spatial_sigma=0.5,
        update_factor=5.0,
        compatibility_kernel_range=1,
        iterations=ITERATIONS,
    )

    runmonaicrf(I_pt, P_pt, crf_func)

    ## output of the script on my local machine:
    # func:'densecrf3d' took: 5.8184 sec
    # func:'runmonaicrf' took: 177.5296 sec

hey @masadcv ,

Thanks you for your help.

It would be grreat , what is your slack ?

Hi @AnisBoubekeur
If you join MONAI slack using instructions from here: https://github.com/Project-MONAI/MONAI#community
You can send me message at: Muhammad Asad with the same profile picture.
I would be happy to assist you further with this setup.