Filtering models and example video clips

Filter2D_Lrec/DataGen.py

@@ -0,0 +1,87 @@
+"""
+A module for generating batches of 2D images.  
+"""
+import os
+import numpy as np
+import random
+import scipy.io as sio
+import tensorflow as tf
+
+def _shift_first_frame(vol_in, vol_out, tr_phase):
+    n_frames = vol_in.shape[-1]
+    n, p = 1, 0.5 # p is the probability of shifting the first frame.
+    if tr_phase:
+        if np.random.binomial(n,p):
+            first_frm = np.random.permutation(np.arange(n_frames))[0]
+            vol_in = np.concatenate((vol_in[:,:,first_frm:], vol_in[:,:,:first_frm]), axis=-1)
+            vol_out = np.concatenate((vol_out[:,:,first_frm:], vol_out[:,:,:first_frm]), axis=-1)        	
+    return vol_in, vol_out
+
+def _image_vol_normalization(vol):
+    vol = vol/255
+    vol[vol < 0] = 0
+    vol[vol > 1] = 1
+    return vol
+
+def _reshape_vol(vol):
+    new_vol = np.empty([vol.shape[-1], vol.shape[0], vol.shape[1]])
+    for i in range(vol.shape[-1]):
+	    new_vol[i,:,:] = vol[:,:,i]
+    return new_vol
+
+def _image_vol_augmentation(vol_in, vol_out, tr_phase):
+    """ 
+    Augmenting and normalizing the input and output image volumes. 
+    """
+    vol_in, vol_out = _shift_first_frame(vol_in, vol_out, tr_phase)
+    vol_in_norm = _image_vol_normalization(vol_in)
+    vol_out_norm = _image_vol_normalization(vol_out)
+    return [_reshape_vol(vol_in_norm), _reshape_vol(vol_out_norm)]
+
+class DataGen(tf.keras.utils.Sequence):
+    """ 
+    Generating batches of input cluttered volumes and their corresponding
+    clutter-free output volumes 
+    """
+    def __init__(
+        self,
+        dim:list,
+        in_dir:str,
+        out_dir:str,
+        id_list:list,
+        batch_size:int,
+        tr_phase=True,
+	    *args,
+        **kwargs):
+        'Initialization'
+        self.dim = dim
+        self.in_dir = in_dir
+        self.out_dir = out_dir
+        self.id_list = id_list
+        self.batch_size = batch_size
+        self.tr_phase = tr_phase
+		
+    def __len__(self):
+        return len(self.id_list)
+    
+    def __getitem__(self, idx):
+        'Generate one or more batches of data'
+        
+        # Initialization
+        in_out_shape = [self.batch_size, self.dim[0], self.dim[1], self.dim[2]]
+        x_aug, y_aug = np.empty(in_out_shape), np.empty(in_out_shape) 
+        vol_id = self.id_list[idx:(idx + 1)]
+        
+        # Generate the data
+        for i, ID in enumerate(vol_id):	
+            # Store sample
+            vol_in = sio.loadmat(self.in_dir[ID])['data_artf']
+            vol_out = sio.loadmat(self.out_dir[ID])['data_org']           
+            # Call the data augmentation function 
+            aug_vols = _image_vol_augmentation(vol_in, vol_out, self.tr_phase)  
+            x_aug[:,:,:,0], y_aug[:,:,:,0] = aug_vols[0], aug_vols[1]
+        
+        if self.tr_phase:
+          return x_aug, y_aug
+        else:
+          return x_aug

Filter2D_Lrec/Error_analysis.py

@@ -0,0 +1,61 @@
+"""
+A module with functions for computing different MAE and coherence errors.
+"""
+import os
+import numpy as np
+import pandas as pd
+import scipy.io as sio
+
+def _compute_sample_temporal_coherency_score(filtered_smp, org_smp):
+    abs_diff_flt_org_smp = np.abs(filtered_smp-org_smp)
+    frm_pixel_sum = [np.sum(abs_diff_flt_org_smp[:,:,i]) for i in range(abs_diff_flt_org_smp.shape[-1])]
+    frm_pixel_sum_shifted = np.roll(frm_pixel_sum, -1)
+    frm_diff = np.abs(frm_pixel_sum - frm_pixel_sum_shifted)
+    frm_diff = frm_diff[:-1]
+    return np.mean(frm_diff)
+
+def _compute_sample_mae(smp, in_ids, filtered_smp, clutter_class):
+    clt_smp = sio.loadmat(in_ids[smp])['data_artf']
+    org_smp = sio.loadmat(
+        in_ids[smp].split(f'data_{clutter_class}')[0] + 'data_org/1.mat')['data_org']
+    mae_CltFiltered_CltFree = np.mean(np.abs(255*filtered_smp-org_smp))
+    mae_Cltrd_CltFree = np.mean(np.abs(clt_smp-org_smp))
+    temporal_coherency_score = _compute_sample_temporal_coherency_score(255*filtered_smp, org_smp)
+    return mae_CltFiltered_CltFree, mae_Cltrd_CltFree, temporal_coherency_score
+
+def _make_res_dct():
+    res_dct = {'Clutter_class': [],
+               'Clutter_spec': [],
+               'View': [],
+               'Vendor': [],
+               'MAE_CltFiltered_CltFree': [],
+               'MAE_Cltrd_CltFree': [],
+               'temporal_coherency_score': []
+               }
+    return res_dct
+
+def _id_separation(in_id):
+    id_part0 = in_id.split('/A')[0].split('/')
+    id_part1 = in_id.split('/data_')[1].split('/')
+    v = [v for v in in_id.split('/') if 'A' in v and 'C' in v]
+    view = v[0]
+    vendor = id_part0[-1]
+    clutter_class = id_part1[0]
+    clutter_spec = id_part1[1]
+    return view, vendor, clutter_class, clutter_spec
+
+def compute_mae(in_ids, filtered_dta, te_subsample=False, te_frames=50):
+    res_dct = _make_res_dct()
+    for i in range(len(in_ids)):
+        view, vendor, clutter_class, clutter_spec = _id_separation(in_ids[i])
+        res_dct['Clutter_class'].append(clutter_class)
+        res_dct['Clutter_spec'].append(clutter_spec)
+        res_dct['Vendor'].append(vendor)
+        res_dct['View'].append(view)
+        mae_CltFiltered_CltFree, mae_Cltrd_CltFree, temporal_coherency_score = _compute_sample_mae(
+            smp=i, in_ids=in_ids, filtered_smp=filtered_dta[i],
+            clutter_class=clutter_class)
+        res_dct['MAE_CltFiltered_CltFree'].append(mae_CltFiltered_CltFree)
+        res_dct['MAE_Cltrd_CltFree'].append(mae_Cltrd_CltFree)
+        res_dct['temporal_coherency_score'].append(temporal_coherency_score)
+    return pd.DataFrame(res_dct)

Filter2D_Lrec/Model_ClutterFilter2D.py

@@ -0,0 +1,128 @@
+"""
+Functions of the 2D clutter filtering algorithm.
+"""
+import tensorflow as tf
+from tensorflow.keras import backend as K
+from tensorflow.keras.layers import (Conv2D, MaxPooling2D, Activation, BatchNormalization, Add,
+                                     Dropout, Concatenate, UpSampling2D, multiply, Input, Lambda) 
+from tensorflow.keras.models import Model
+from tensorflow.keras.optimizers import Adam
+
+def tensor_expansion(tensor, rep, axs):
+    expanded_tensor = Lambda(lambda x, repnum: K.repeat_elements(x, repnum, axis=axs),
+                             arguments={'repnum': rep})(tensor)
+    return expanded_tensor
+
+def attention_gate_block_2D(x, g, n_inter_filters=None, name=None, **prm):
+    """ 
+    Attention gate block.
+    """
+    shape_x = K.int_shape(x)
+    shape_g = K.int_shape(g) 
+    if n_inter_filters is None:
+       n_inter_filters = shape_x[-1] // 2
+         
+    theta_x = Conv2D(n_inter_filters, 3, strides=(2, 2), padding='same', name=f"{name}_theta_x")(x)
+    phi_g = Conv2D(n_inter_filters, 1, strides=1, padding='valid', name=f"{name}_phi_g")(g) 
+    concat_xg = Add()([phi_g, theta_x])
+    act_xg = Activation('relu')(concat_xg)
+    psi = Conv2D(1, 1, padding='same', name=f"{name}_psi")(act_xg)
+    sigmoid_xg = Activation('sigmoid')(psi)
+    shape_sigmoid = K.int_shape(sigmoid_xg)
+    upsample_sigmoid = UpSampling2D(size=(2, 2), name=f"{name}_upsampled_sig")(sigmoid_xg)
+    upsample_sigmoid_rep = tensor_expansion(upsample_sigmoid, rep=shape_x[-1], axs=-1)
+    y = multiply([upsample_sigmoid_rep, x], name=f"{name}_weighted_x")
+    return y
+
+def conv_block(x, filters_list, act='linear', kernel_size=3, stride=1, pad='same', drp=0.05, name=None):
+    """ 
+    Blocks of 2D conv filters.
+    """
+    for i in range(len(filters_list)):	
+        x = Conv2D(filters_list[i], kernel_size, padding=pad, strides=stride, name=f"{name}_blk{i+1}")(x)
+        x = BatchNormalization(name=f"{name}_bn{i+1}")(x)
+        x = Activation(act, name=f"{name}_act{i+1}")(x)
+        x = Dropout(drp)(x)
+    return x
+
+def encoding_block(x_in, name, **config):
+    """
+    Encoding block of the 2D Unet.
+    """
+    encoding_dct = {}
+    for i in range(config["network_prm"]["n_levels"]):        
+        if i == 0:
+            x = x_in  
+            n_filters = config["network_prm"]["n_init_filters"]
+        else:
+            n_filters = (2**i)*config["network_prm"]["n_init_filters"]
+            x = MaxPooling2D(pool_size=config["network_prm"]["pool_size"], name=f"{name}_encd_pool{i}")(x)
+        x = conv_block(x, filters_list=[n_filters, 2*n_filters], act=config["network_prm"]["act"],
+	                   kernel_size=config["network_prm"]["kernel_size"],
+                       stride=config["network_prm"]["conv_stride"], 
+                       pad=config["network_prm"]["padding"],
+                       drp=config["learning_prm"]['drp'], name=f"{name}_encd_conv_lvl{i}")
+        encoding_dct[f"{name}_out_lvl{i}"] = x
+    return encoding_dct
+
+def decoding_block(encoding_dct, name, **config):
+    """
+    Decoding block of the 2D Unet.
+    """
+    decoding_dct = {}
+    n_levels = config["network_prm"]["n_levels"]
+    for i in range(n_levels-1):
+        if i == 0:
+            x = encoding_dct[f"{name}_out_lvl{n_levels-i-1}"]
+	    # upsampling via Conv(Upsampling)
+        x_shape = K.int_shape(x)
+        x_up = Conv2D(x_shape[-1], 2, activation=config["network_prm"]["act"], padding='same', strides=1,
+                      name=f"{name}_decd_upsmpl{i}")(UpSampling2D(size=(2,2))(x))
+        x_up_shape = K.int_shape(x_up)
+        # concatenation
+        if config["network_prm"]['attention']:
+            if i == 0:
+                g = encoding_dct[f"{name}_out_lvl{n_levels-1}"]
+            else:
+                g = decoding_dct[f"{name}_out_lvl{i-1}"]
+            x_encd = attention_gate_block_2D(x=encoding_dct[f"{name}_out_lvl{n_levels-i-2}"], 
+                                             g=g, name=f"{name}_att_blk{i}")
+        else:
+            x_encd = encoding_dct[f"{name}_out_lvl{n_levels-i-2}"]
+        x_concat = Concatenate(axis=-1, name=f"{name}_decd_concat{i}")([x_encd, x_up])
+        n_filters = x_up_shape[-1]//2
+        x = conv_block(x_concat, filters_list=[n_filters, n_filters], act=config["network_prm"]["act"],
+		               kernel_size=config["network_prm"]["kernel_size"],
+                       stride=config["network_prm"]["conv_stride"],
+		               pad=config["network_prm"]["padding"],
+                       drp=config["learning_prm"]['drp'], name=f"{name}_decd_conv_lvl{i}")
+        decoding_dct[f"{name}_out_lvl{i}"] = x
+    x = conv_block(x, filters_list=[1], act=config["network_prm"]["act"], kernel_size=1, stride=1, 
+		           pad='same', drp=1e-4, name=f"{name}_final_decd_conv")
+    decoding_dct[f"{name}_final_conv"] = x
+    return decoding_dct
+
+def Unet2D(x_in, name, **config):
+    """
+    Spatial clutter filtering based on 2D Unet.
+    """
+    encoding_dct = encoding_block(x_in, name, **config)
+    decoding_dct = decoding_block(encoding_dct, name, **config)
+    # Add the Main input here
+    if config["network_prm"]["in_skip"]:
+        out_Unet = Add()([x_in, decoding_dct[f"{name}_final_conv"]])
+    else:
+        out_Unet = decoding_dct[f"{name}_final_conv"]
+    return out_Unet
+
+def clutter_filter_2D(**config):
+    """
+    The main function for designing the clutter filtering algorithm.
+    """
+    main_in = Input(config["network_prm"]["input_dim"])
+    filter_out = Unet2D(x_in=main_in, name="CF", **config)
+    model = Model(inputs=main_in, outputs=filter_out, name=config['model_name'])
+    opt = Adam(learning_rate=config["learning_prm"]["lr"])
+    model.compile(optimizer=opt, loss=config["learning_prm"]["loss"], metrics=config["learning_prm"]["metrics"])
+    model.summary()
+    return model

Filter2D_Lrec/TestClutterFilter2D.py

@@ -0,0 +1,45 @@
+""" 
+Module for testing the 3D clutter filtering model.
+"""
+import os
+import argparse
+import json
+import numpy as np
+import pandas as pd
+
+from utils import *
+from Model_ClutterFilter2D import clutter_filter_2D
+from DataGen import DataGen
+from Error_analysis import compute_mae
+    
+def data_generation(in_ids_te, out_ids_te, config):
+    DtaGenTe_prm = {
+        'dim': config["network_prm"]["input_dim"],
+        'in_dir': in_ids_te,
+        'out_dir': out_ids_te,
+        'id_list': np.arange(len(in_ids_te)),
+        'batch_size': config["learning_prm"]["batch_size"],
+        'tr_phase': False} 
+    return DataGen(**DtaGenTe_prm)
+
+def main(config):
+    in_ids_te, out_ids_te, te_subject, val_subject = id_preparation(config)
+    te_gen = data_generation(in_ids_te, out_ids_te, config)
+    model = clutter_filter_2D(**config)
+    weight_dir = create_weight_dir(val_subject, te_subject, config)
+    model.load_weights(
+        os.path.join(weight_dir, config["weight_name"] + ".hdf5"))
+    results_te = model.predict_generator(te_gen, verbose=2)
+    df_errors = compute_mae(in_ids_te, results_te)
+    df_errors.to_csv(
+        os.path.join(weight_dir, config["weight_name"] + ".csv"))
+    return None
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", help="path of the config file", default="config.json")
+    args = parser.parse_args()
+    assert os.path.isfile(args.config)
+    with open(args.config, "r") as read_file:
+        config = json.load(read_file)
+    main(config)

Filter2D_Lrec/TrainClutterFilter2D.py

@@ -0,0 +1,69 @@
+""" 
+Module for training the 2D clutter filtering model with L2 loss.
+"""
+import os
+import argparse
+import json
+import numpy as np
+from tensorflow.keras.callbacks import ModelCheckpoint, ReduceLROnPlateau
+
+from utils import *
+from Model_ClutterFilter2D import clutter_filter_2D
+from DataGen import DataGen
+
+def data_generation(in_ids_tr, in_ids_val, out_ids_tr, out_ids_val, config):
+    DtaGenTr_prm = {
+        'dim': config["network_prm"]["input_dim"],
+        'in_dir': in_ids_tr,
+        'out_dir': out_ids_tr,
+        'id_list': np.arange(len(in_ids_tr)),
+        'batch_size': config["learning_prm"]["batch_size"],
+        'tr_phase': True}
+    DtaGenVal_prm = {
+        'dim': config["network_prm"]["input_dim"],
+        'in_dir': in_ids_val,
+        'out_dir': out_ids_val,
+        'id_list': np.arange(len(in_ids_val)),
+        'batch_size': config["learning_prm"]["batch_size"],
+        'tr_phase': True}
+    tr_gen = DataGen(**DtaGenTr_prm)
+    val_gen = DataGen(**DtaGenVal_prm)
+    return tr_gen, val_gen
+
+def model_chkpnt(val_subject, te_subject, weight_dir, config):
+    weight_name = (
+            f'CF2D_ValTeSbj_{val_subject}_{te_subject}_nLvl{config["network_prm"]["n_levels"]}'
+            f'_InSkp{config["network_prm"]["in_skip"]}_Att{config["network_prm"]["attention"]}'
+            f'_Act{config["network_prm"]["act"]}_nInitFlt{config["network_prm"]["n_init_filters"]}_lr{config["learning_prm"]["lr"]}')
+    filepath = (weight_dir + '/'+  weight_name +
+                '_epc' + "{epoch:03d}" + '_trloss' + "{loss:.5f}" +
+                '_valloss' + "{val_loss:.5f}" + ".hdf5")
+    model_checkpoint = ModelCheckpoint(filepath=filepath,
+                                       monitor="val_loss",
+                                       verbose=0,
+                                       save_best_only=True)
+    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1,
+                                  patience=4, min_lr=1e-7)
+    return model_checkpoint, reduce_lr
+
+def main(config):
+    in_ids_tr, in_ids_val, out_ids_tr, out_ids_val, val_subject, te_subject = id_preparation(config)
+    weight_dir = create_weight_dir(val_subject, te_subject, config)
+    tr_gen, val_gen = data_generation(in_ids_tr, in_ids_val, out_ids_tr, out_ids_val, config)
+    model = clutter_filter_2D(**config)
+    model_checkpoint, reduce_lr = model_chkpnt(val_subject, te_subject, weight_dir, config)
+    model.fit(tr_gen,
+              validation_data=val_gen,
+              epochs=config["learning_prm"]["n_epochs"],
+              verbose=1,
+              callbacks=[model_checkpoint, reduce_lr])
+    return None
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", help="path of the config file", default="config.json")
+    args = parser.parse_args()
+    assert os.path.isfile(args.config)
+    with open(args.config, "r") as read_file:
+        config = json.load(read_file)
+    main(config)

Filter2D_Lrec/config.json

@@ -0,0 +1,35 @@
+{
+    "paths": {
+        "data_path": "",
+        "save_path": ""
+    },
+    "subject_list": ["rca", "ladprox", "laddist", "lcx", "normal"],
+    "CV": {
+        "val_subject_id": 0
+    },
+    "network_prm": {
+        "input_dim": [128, 128, 1],
+        "n_levels": 4,
+        "n_init_filters": 16,
+        "in_skip": true,
+        "attention": true,
+        "kernel_size": 3,
+        "conv_stride": 1,
+        "upsampling_stride": [2, 2],
+        "pool_size": [2, 2], 
+        "pool_stride": 1,
+        "padding": "same",
+        "act": "linear"
+    },
+    "learning_prm": {
+        "batch_size": 50,
+        "lr": 1e-4,
+        "drp": 0.05,
+        "loss": "mean_squared_error",
+        "metrics": ["mae"],
+        "n_epochs": 10
+    },
+    "tr_phase": true,
+    "model_name": "CF2D_L2Loss",
+    "weight_name": ""
+}

Filter2D_Lrec/utils.py

@@ -0,0 +1,56 @@
+""" 
+Utility functions.
+"""
+import os
+import numpy as np
+
+def generate_tr_val_te_subject_ids(subject_list, val_subject_id):
+	val_subject = subject_list[val_subject_id]
+	te_subject = subject_list[val_subject_id-1]
+	subject_list.remove(val_subject)
+	subject_list.remove(te_subject)
+	tr_subjects = subject_list
+	return tr_subjects, val_subject, te_subject
+
+def generate_data_ids(data_dir, subject_list):
+    in_ids, out_ids = [], []
+    vendor_list = [vendor for vendor in os.listdir(data_dir) if '.' not in vendor]
+    for vendor in vendor_list:
+        vendor_dir = os.path.join(data_dir, vendor)
+        view_list = [view for view in os.listdir(vendor_dir) if '.' not in view]
+        for view in view_list:
+            view_dir = os.path.join(vendor_dir, view) 
+            subject_full_list = [subject for subject in os.listdir(view_dir) if '.' not in subject]
+            for subject in subject_full_list:
+                if subject in subject_list:
+                    subject_dir = os.path.join(view_dir, subject)
+                    org_data_dir = os.path.join(subject_dir, 'data_org')
+                    org_data_id = os.path.join(org_data_dir, os.listdir(org_data_dir)[0])
+                    clutter_list = [clutter for clutter in os.listdir(subject_dir)
+                                    if clutter in ['data_NFClt', 'data_NFRvbClt', 'data_RvbClt']
+                                    and '.' not in clutter]
+                    for clutter in clutter_list:
+                        clutter_dir = os.path.join(subject_dir, clutter)
+                        clutter_ids = os.listdir(clutter_dir)
+                        clutter_ids_dir = [os.path.join(clutter_dir, id_dir) for id_dir in clutter_ids if '.DS' not in id_dir]
+                        in_ids += clutter_ids_dir
+                        out_ids += [org_data_id]*len(os.listdir(clutter_dir))
+    return in_ids, out_ids
+
+def id_preparation(config):
+    tr_subjects, val_subject, te_subject = generate_tr_val_te_subject_ids(
+        subject_list=config["subject_list"], val_subject_id=config["CV"]["val_subject_id"])
+    if config["tr_phase"]:
+        in_ids_tr, out_ids_tr = generate_data_ids(config["paths"]["data_path"], tr_subjects)
+        in_ids_val, out_ids_val = generate_data_ids(config["paths"]["data_path"], val_subject)
+        return in_ids_tr, in_ids_val, out_ids_tr, out_ids_val, val_subject, te_subject
+    else:
+        in_ids_te, out_ids_te = generate_data_ids(config["paths"]["data_path"], te_subject)
+        return in_ids_te, out_ids_te, te_subject, val_subject
+
+def create_weight_dir(val_subject, te_subject, config):
+    weight_dir = os.path.join(config["paths"]["save_path"],
+                              "Weights", f"ValTeIDs_{val_subject}_{te_subject}")  
+    if not os.path.exists(weight_dir):
+        os.makedirs(weight_dir)
+    return weight_dir

Filter3D_Lrec/DataGen.py

@@ -0,0 +1,77 @@
+"""
+Module for generating batches of 2D images.  
+"""
+import os
+import numpy as np
+import random
+import scipy.io as sio
+import tensorflow as tf
+
+def _shift_first_frame(vol_in, vol_out, tr_phase):
+    n_frames = vol_in.shape[-1]
+    n, p = 1, 0.5 # p is the probability of shifting the first frame.
+    if tr_phase:
+        if np.random.binomial(n,p):
+            first_frm = np.random.permutation(np.arange(n_frames))[0]
+            vol_in = np.concatenate((vol_in[:,:,first_frm:], vol_in[:,:,:first_frm]), axis=-1)
+            vol_out = np.concatenate((vol_out[:,:,first_frm:], vol_out[:,:,:first_frm]), axis=-1)        	
+    return vol_in, vol_out
+
+def _image_vol_normalization(vol):
+    vol = vol/255
+    vol[vol < 0] = 0
+    vol[vol > 1] = 1
+    return vol
+
+def _image_vol_augmentation(vol_in, vol_out, tr_phase):
+    vol_in, vol_out = _shift_first_frame(vol_in, vol_out, tr_phase)
+    vol_in_norm = _image_vol_normalization(vol_in)
+    vol_out_norm = _image_vol_normalization(vol_out)
+    vol_shape = [sh for sh in vol_in.shape]
+    vol_shape.append(1)
+    vol_in, vol_out = np.empty(vol_shape), np.empty(vol_shape)
+    vol_in[:,:,:,0], vol_out[:,:,:,0] = vol_in_norm, vol_out_norm
+    return [vol_in, vol_out]
+
+class DataGen(tf.keras.utils.Sequence):
+    """ 
+    Generating batches of input cluttered volumes and their corresponding
+    clutter-free output volumes 
+    """
+    def __init__(
+        self,
+        dim:list,
+        in_dir:str,
+        out_dir:str,
+        id_list:list,
+        batch_size:int,
+        tr_phase=True,
+        te_subsample=False,
+        te_frames=0,
+	    *args,
+        **kwargs):
+        'Initialization'
+        self.dim = dim
+        self.in_dir = in_dir
+        self.out_dir = out_dir
+        self.id_list = id_list
+        self.batch_size = batch_size
+        self.tr_phase = tr_phase
+		
+    def __len__(self):
+        return int(np.floor(len(self.id_list) / self.batch_size))
+    
+    def __getitem__(self, idx):
+        batch = self.id_list[idx*self.batch_size:(idx+1)*self.batch_size]
+        x_aug, y_aug = [], []
+        for i, ID in enumerate(batch):
+            vol_in = sio.loadmat(self.in_dir[ID])['data_artf']
+            vol_out = sio.loadmat(self.out_dir[ID])['data_org']           
+            # Call the data augmentation function 
+            aug_vols = _image_vol_augmentation(vol_in, vol_out, self.tr_phase)
+            x_aug.append(aug_vols[0])
+            y_aug.append(aug_vols[1])
+        if self.tr_phase:
+            return np.asarray(x_aug), np.asarray(y_aug)
+        else:
+            return np.asarray(x_aug)

Filter3D_Lrec/Error_analysis.py

@@ -0,0 +1,61 @@
+"""
+A module with functions for computing different MAE and coherence errors.  
+"""
+import os
+import numpy as np
+import pandas as pd
+import scipy.io as sio
+
+def _compute_sample_temporal_coherency_score(filtered_smp, org_smp):
+    abs_diff_flt_org_smp = np.abs(filtered_smp-org_smp)
+    frm_pixel_sum = [np.sum(abs_diff_flt_org_smp[:,:,i]) for i in range(abs_diff_flt_org_smp.shape[-1])]
+    frm_pixel_sum_shifted = np.roll(frm_pixel_sum, -1)
+    frm_diff = np.abs(frm_pixel_sum - frm_pixel_sum_shifted)
+    frm_diff = frm_diff[:-1]
+    return np.mean(frm_diff)
+
+def _compute_sample_mae(smp, in_ids, filtered_smp, clutter_class):
+    clt_smp = sio.loadmat(in_ids[smp])['data_artf']
+    org_smp = sio.loadmat(
+        in_ids[smp].split(f'data_{clutter_class}')[0] + 'data_org/1.mat')['data_org']
+    mae_CltFiltered_CltFree = np.mean(np.abs(255*filtered_smp-org_smp))
+    mae_Cltrd_CltFree = np.mean(np.abs(clt_smp-org_smp))
+    temporal_coherency_score = _compute_sample_temporal_coherency_score(255*filtered_smp, org_smp)
+    return mae_CltFiltered_CltFree, mae_Cltrd_CltFree, temporal_coherency_score
+
+def _make_res_dct():
+    res_dct = {'Clutter_class': [],
+               'Clutter_spec': [],
+               'View': [],
+               'Vendor': [],
+               'MAE_CltFiltered_CltFree': [],
+               'MAE_Cltrd_CltFree': [],
+               'temporal_coherency_score': []
+               }
+    return res_dct
+
+def _id_separation(in_id):
+    id_part0 = in_id.split('/A')[0].split('/')
+    id_part1 = in_id.split('/data_')[1].split('/')
+    v = [v for v in in_id.split('/') if 'A' in v and 'C' in v]
+    view = v[0]
+    vendor = id_part0[-1]
+    clutter_class = id_part1[0]
+    clutter_spec = id_part1[1]
+    return view, vendor, clutter_class, clutter_spec
+
+def compute_mae(in_ids, filtered_dta, te_subsample=False, te_frames=50):
+    res_dct = _make_res_dct()
+    for i in range(len(in_ids)):
+        view, vendor, clutter_class, clutter_spec = _id_separation(in_ids[i])
+        res_dct['Clutter_class'].append(clutter_class)
+        res_dct['Clutter_spec'].append(clutter_spec)
+        res_dct['Vendor'].append(vendor)
+        res_dct['View'].append(view)
+        mae_CltFiltered_CltFree, mae_Cltrd_CltFree, temporal_coherency_score = _compute_sample_mae(
+            smp=i, in_ids=in_ids, filtered_smp=filtered_dta[i,:,:,:,0],
+            clutter_class=clutter_class)
+        res_dct['MAE_CltFiltered_CltFree'].append(mae_CltFiltered_CltFree)
+        res_dct['MAE_Cltrd_CltFree'].append(mae_Cltrd_CltFree)
+        res_dct['temporal_coherency_score'].append(temporal_coherency_score)
+    return pd.DataFrame(res_dct)

Filter3D_Lrec/Model_ClutterFilter3D.py

@@ -0,0 +1,128 @@
+"""
+Functions of the 3D clutter filtering algorithm.
+"""
+import tensorflow as tf
+from tensorflow.keras import backend as K
+from tensorflow.keras.layers import (Conv3D, MaxPooling3D, Activation, BatchNormalization, Add,
+                                     Dropout, Concatenate, UpSampling3D, multiply, Input, Lambda) 
+from tensorflow.keras.models import Model
+from tensorflow.keras.optimizers import Adam
+
+def tensor_expansion(tensor, rep, axs):
+    expanded_tensor = Lambda(lambda x, repnum: K.repeat_elements(x, repnum, axis=axs),
+                             arguments={'repnum': rep})(tensor)
+    return expanded_tensor
+
+def attention_gate_block_3D(x, g, n_inter_filters=None, name=None, **prm):
+    """ 
+    Attention gate block. 
+    """
+    shape_x = K.int_shape(x)
+    shape_g = K.int_shape(g) 
+    if n_inter_filters is None:
+       n_inter_filters = shape_x[-1] // 2
+         
+    theta_x = Conv3D(n_inter_filters, 3, strides=(2, 2, 1), padding='same', name=f"{name}_theta_x")(x)
+    phi_g = Conv3D(n_inter_filters, 1, strides=1, padding='valid', name=f"{name}_phi_g")(g) 
+    concat_xg = Add()([phi_g, theta_x])
+    act_xg = Activation('relu')(concat_xg)
+    psi = Conv3D(1, 1, padding='same', name=f"{name}_psi")(act_xg)
+    sigmoid_xg = Activation('sigmoid')(psi)
+    shape_sigmoid = K.int_shape(sigmoid_xg)
+    upsample_sigmoid = UpSampling3D(size=(2, 2, 1), name=f"{name}_upsampled_sig")(sigmoid_xg)
+    upsample_sigmoid_rep = tensor_expansion(upsample_sigmoid, rep=shape_x[-1], axs=-1)
+    y = multiply([upsample_sigmoid_rep, x], name=f"{name}_weighted_x")
+    return y
+
+def conv_block(x, filters_list, act='linear', kernel_size=3, stride=1, pad='same', drp=0.05, name=None):
+    """ 
+    Blocks of 3D conv filters.
+    """
+    for i in range(len(filters_list)):	
+        x = Conv3D(filters_list[i], kernel_size, padding=pad, strides=stride, name=f"{name}_blk{i+1}")(x)
+        x = BatchNormalization(name=f"{name}_bn{i+1}")(x)
+        x = Activation(act, name=f"{name}_act{i+1}")(x)
+        x = Dropout(drp)(x)
+    return x
+
+def encoding_block(x_in, name, **config):
+    """
+    Encoding block of the 3D Unet.
+    """
+    encoding_dct = {}
+    for i in range(config["network_prm"]["n_levels"]):        
+        if i == 0:
+            x = x_in  
+            n_filters = config["network_prm"]["n_init_filters"]
+        else:
+            n_filters = (2**i)*config["network_prm"]["n_init_filters"]
+            x = MaxPooling3D(pool_size=config["network_prm"]["pool_size"], name=f"{name}_encd_pool{i}")(x)
+        x = conv_block(x, filters_list=[n_filters, 2*n_filters], act=config["network_prm"]["act"],
+	                   kernel_size=config["network_prm"]["kernel_size"],
+                       stride=config["network_prm"]["conv_stride"], 
+                       pad=config["network_prm"]["padding"],
+                       drp=config["learning_prm"]['drp'], name=f"{name}_encd_conv_lvl{i}")
+        encoding_dct[f"{name}_out_lvl{i}"] = x
+    return encoding_dct
+
+def decoding_block(encoding_dct, name, **config):
+    """
+    Decoding block of the 3D Unet.
+    """
+    decoding_dct = {}
+    n_levels = config["network_prm"]["n_levels"]
+    for i in range(n_levels-1):
+        if i == 0:
+            x = encoding_dct[f"{name}_out_lvl{n_levels-i-1}"]
+	    # upsampling via Conv(Upsampling)
+        x_shape = K.int_shape(x)
+        x_up = Conv3D(x_shape[-1], 2, activation=config["network_prm"]["act"], padding='same', strides=1,
+                      name=f"{name}_decd_upsmpl{i}")(UpSampling3D(size=(2,2,1))(x))
+        x_up_shape = K.int_shape(x_up)
+        # concatenation
+        if config["network_prm"]['attention']:
+            if i == 0:
+                g = encoding_dct[f"{name}_out_lvl{n_levels-1}"]
+            else:
+                g = decoding_dct[f"{name}_out_lvl{i-1}"]
+            x_encd = attention_gate_block_3D(x=encoding_dct[f"{name}_out_lvl{n_levels-i-2}"], 
+                                             g=g, name=f"{name}_att_blk{i}")
+        else:
+            x_encd = encoding_dct[f"{name}_out_lvl{n_levels-i-2}"]
+        x_concat = Concatenate(axis=-1, name=f"{name}_decd_concat{i}")([x_encd, x_up])
+        n_filters = x_up_shape[-1]//2
+        x = conv_block(x_concat, filters_list=[n_filters, n_filters], act=config["network_prm"]["act"],
+		               kernel_size=config["network_prm"]["kernel_size"],
+                       stride=config["network_prm"]["conv_stride"],
+		               pad=config["network_prm"]["padding"],
+                       drp=config["learning_prm"]['drp'], name=f"{name}_decd_conv_lvl{i}")
+        decoding_dct[f"{name}_out_lvl{i}"] = x
+    x = conv_block(x, filters_list=[1], act=config["network_prm"]["act"], kernel_size=1, stride=1, 
+		           pad='same', drp=1e-4, name=f"{name}_final_decd_conv")
+    decoding_dct[f"{name}_final_conv"] = x
+    return decoding_dct
+
+def Unet3D(x_in, name, **config):
+    """
+    Spatiotemporal clutter filtering model based on the 3D Unet.
+    """
+    encoding_dct = encoding_block(x_in, name, **config)
+    decoding_dct = decoding_block(encoding_dct, name, **config)
+    if config["network_prm"]["in_skip"]:
+        out_Unet = Add()([x_in, decoding_dct[f"{name}_final_conv"]])
+    else:
+        out_Unet = decoding_dct[f"{name}_final_conv"]
+    return out_Unet
+
+def clutter_filter_3D(**config):
+    """
+    The main function for designing the clutter filtering algorithm.
+    """
+    main_in = Input(config["network_prm"]["input_dim"])
+    filter_out = Unet3D(x_in=main_in, name="CF", **config)
+    model = Model(inputs=main_in, outputs=filter_out, name=config['model_name'])
+    opt = Adam(learning_rate=config["learning_prm"]["lr"])
+    model.compile(optimizer=opt, loss=config["learning_prm"]["loss"],
+                  metrics=config["learning_prm"]["metrics"])
+    model.summary()
+    return model

Filter3D_Lrec/TestClutterFilter3D.py

@@ -0,0 +1,45 @@
+""" 
+Module for testing the 3D clutter filtering model.
+"""
+import os
+import argparse
+import json
+import numpy as np
+import pandas as pd
+
+from utils import *
+from Model_ClutterFilter3D import clutter_filter_3D
+from DataGen import DataGen
+from Error_analysis import compute_mae
+    
+def data_generation(in_ids_te, out_ids_te, config):
+    DtaGenTe_prm = {
+        'dim': config["network_prm"]["input_dim"],
+        'in_dir': in_ids_te,
+        'out_dir': out_ids_te,
+        'id_list': np.arange(len(in_ids_te)),
+        'batch_size': config["learning_prm"]["batch_size"],
+        'tr_phase': False} 
+    return DataGen(**DtaGenTe_prm)
+
+def main(config):
+    in_ids_te, out_ids_te, te_subject, val_subject = id_preparation(config)
+    te_gen = data_generation(in_ids_te, out_ids_te, config)
+    model = clutter_filter_3D(**config)
+    weight_dir = create_weight_dir(val_subject, te_subject, config)
+    model.load_weights(
+        os.path.join(weight_dir, config["weight_name"] + ".hdf5"))
+    results_te = model.predict_generator(te_gen, verbose=2)
+    df_errors = compute_mae(in_ids_te, results_te)
+    df_errors.to_csv(
+        os.path.join(weight_dir, config["weight_name"] + ".csv"))
+    return None
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", help="path of the config file", default="config.json")
+    args = parser.parse_args()
+    assert os.path.isfile(args.config)
+    with open(args.config, "r") as read_file:
+        config = json.load(read_file)
+    main(config)

Filter3D_Lrec/TrainClutterFilter3D.py

@@ -0,0 +1,70 @@
+""" 
+Module for training the 3D clutter filtering model with L2 loss. 
+"""
+import os
+import argparse
+import json
+import numpy as np
+from tensorflow.keras.callbacks import ModelCheckpoint, ReduceLROnPlateau
+
+from utils import *
+from Model_ClutterFilter3D import clutter_filter_3D
+from DataGen import DataGen
+
+def data_generation(in_ids_tr, in_ids_val, out_ids_tr, out_ids_val, config):
+    DtaGenTr_prm = {
+        'dim': config["network_prm"]["input_dim"],
+        'in_dir': in_ids_tr,
+        'out_dir': out_ids_tr,
+        'id_list': np.arange(len(in_ids_tr)),
+        'batch_size': config["learning_prm"]["batch_size"],
+        'tr_phase': True}
+    DtaGenVal_prm = {
+        'dim': config["network_prm"]["input_dim"],
+        'in_dir': in_ids_val,
+        'out_dir': out_ids_val,
+        'id_list': np.arange(len(in_ids_val)),
+        'batch_size': config["learning_prm"]["batch_size"],
+        'tr_phase': True}
+    tr_gen = DataGen(**DtaGenTr_prm)
+    val_gen = DataGen(**DtaGenVal_prm)
+    return tr_gen, val_gen
+
+def model_chkpnt(val_subject, te_subject, weight_dir, config):
+    weight_name = (
+            f'CF3D_ValTeSbj_{val_subject}_{te_subject}_nLvl{config["network_prm"]["n_levels"]}'
+            f'_InSkp{config["network_prm"]["in_skip"]}_Att{config["network_prm"]["attention"]}'
+            f'_Act{config["network_prm"]["act"]}_nInitFlt{config["network_prm"]["n_init_filters"]}'
+            f'_lr{config["learning_prm"]["lr"]}')
+    filepath = (weight_dir + '/'+  weight_name +
+                '_epc' + "{epoch:03d}" + '_trloss' + "{loss:.5f}" +
+                '_valloss' + "{val_loss:.5f}" + ".hdf5")
+    model_checkpoint = ModelCheckpoint(filepath=filepath,
+                                       monitor="val_loss",
+                                       verbose=0,
+                                       save_best_only=True)
+    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1,
+                                  patience=4, min_lr=1e-7)
+    return model_checkpoint, reduce_lr
+
+def main(config):
+    in_ids_tr, in_ids_val, out_ids_tr, out_ids_val, val_subject, te_subject = id_preparation(config)
+    weight_dir = create_weight_dir(val_subject, te_subject, config)
+    tr_gen, val_gen = data_generation(in_ids_tr, in_ids_val, out_ids_tr, out_ids_val, config)
+    model = clutter_filter_3D(**config)
+    model_checkpoint, reduce_lr = model_chkpnt(val_subject, te_subject, weight_dir, config)
+    model.fit(tr_gen,
+              validation_data=val_gen,
+              epochs=config["learning_prm"]["n_epochs"],
+              verbose=1,
+              callbacks=[model_checkpoint, reduce_lr])
+    return None
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", help="path of the config file", default="config.json")
+    args = parser.parse_args()
+    assert os.path.isfile(args.config)
+    with open(args.config, "r") as read_file:
+        config = json.load(read_file)
+    main(config)

Filter3D_Lrec/config.json

@@ -0,0 +1,35 @@
+{
+    "paths": {
+        "data_path": "",
+        "save_path": ""
+    },
+    "subject_list": ["rca", "ladprox", "laddist", "lcx", "normal"],
+    "CV": {
+        "val_subject_id": 0
+    },
+    "network_prm": {
+        "input_dim": [128, 128, 50, 1],
+        "n_levels": 4,
+        "n_init_filters": 16,
+        "in_skip": true,
+        "attention": true,
+        "kernel_size": 3,
+        "conv_stride": 1,
+        "upsampling_stride": [2, 2, 1],
+        "pool_size": [2, 2, 1], 
+        "pool_stride": 1,
+        "padding": "same",
+        "act": "linear"
+    },
+    "learning_prm": {
+        "batch_size": 1,
+        "lr": 1e-4,
+        "drp": 0.05,
+        "loss": "mean_squared_error",
+        "metrics": ["mae"],
+        "n_epochs": 10
+    },
+    "tr_phase": true,
+    "model_name": "CF3D_L2Loss",
+    "weight_name": ""
+}

Filter3D_Lrec/utils.py

@@ -0,0 +1,56 @@
+""" 
+Utility functions.
+"""
+import os
+import numpy as np
+
+def generate_tr_val_te_subject_ids(subject_list, val_subject_id):
+	val_subject = subject_list[val_subject_id]
+	te_subject = subject_list[val_subject_id-1]
+	subject_list.remove(val_subject)
+	subject_list.remove(te_subject)
+	tr_subjects = subject_list
+	return tr_subjects, val_subject, te_subject
+
+def generate_data_ids(data_dir, subject_list):
+    in_ids, out_ids = [], []
+    vendor_list = [vendor for vendor in os.listdir(data_dir) if '.' not in vendor]
+    for vendor in vendor_list:
+        vendor_dir = os.path.join(data_dir, vendor)
+        view_list = [view for view in os.listdir(vendor_dir) if '.' not in view]
+        for view in view_list:
+            view_dir = os.path.join(vendor_dir, view) 
+            subject_full_list = [subject for subject in os.listdir(view_dir) if '.' not in subject]
+            for subject in subject_full_list:
+                if subject in subject_list:
+                    subject_dir = os.path.join(view_dir, subject)
+                    org_data_dir = os.path.join(subject_dir, 'data_org')
+                    org_data_id = os.path.join(org_data_dir, os.listdir(org_data_dir)[0])
+                    clutter_list = [clutter for clutter in os.listdir(subject_dir)
+                                    if clutter in ['data_NFClt', 'data_NFRvbClt', 'data_RvbClt']
+                                    and '.' not in clutter]
+                    for clutter in clutter_list:
+                        clutter_dir = os.path.join(subject_dir, clutter)
+                        clutter_ids = os.listdir(clutter_dir)
+                        clutter_ids_dir = [os.path.join(clutter_dir, id_dir) for id_dir in clutter_ids if '.DS' not in id_dir]
+                        in_ids += clutter_ids_dir
+                        out_ids += [org_data_id]*len(os.listdir(clutter_dir))
+    return in_ids, out_ids
+
+def id_preparation(config):
+    tr_subjects, val_subject, te_subject = generate_tr_val_te_subject_ids(
+        subject_list=config["subject_list"], val_subject_id=config["CV"]["val_subject_id"])
+    if config["tr_phase"]:
+        in_ids_tr, out_ids_tr = generate_data_ids(config["paths"]["data_path"], tr_subjects)
+        in_ids_val, out_ids_val = generate_data_ids(config["paths"]["data_path"], val_subject)
+        return in_ids_tr, in_ids_val, out_ids_tr, out_ids_val, val_subject, te_subject
+    else:
+        in_ids_te, out_ids_te = generate_data_ids(config["paths"]["data_path"], te_subject)
+        return in_ids_te, out_ids_te, te_subject, val_subject
+
+def create_weight_dir(val_subject, te_subject, config):
+    weight_dir = os.path.join(config["paths"]["save_path"],
+                              "Weights", f"ValTeIDs_{val_subject}_{te_subject}")  
+    if not os.path.exists(weight_dir):
+        os.makedirs(weight_dir)
+    return weight_dir

Filter3D_Lrec_adv/DataGen.py

@@ -0,0 +1,111 @@
+"""
+Module for generating batches of 3D images and a 3D mask.  
+"""
+import os
+import numpy as np
+import random
+import scipy.io as sio
+import tensorflow as tf
+ 
+def _shift_first_frame(vol_in, vol_out, tr_phase):
+    n_frames = vol_in.shape[-1]
+    n, p = 1, 0.5 # p is the probability of shifting the first frame.
+    if tr_phase:
+        if np.random.binomial(n,p):
+            first_frm = np.random.permutation(np.arange(n_frames))[0]
+            vol_in = np.concatenate((vol_in[:,:,first_frm:], vol_in[:,:,:first_frm]), axis=-1)
+            vol_out = np.concatenate((vol_out[:,:,first_frm:], vol_out[:,:,:first_frm]), axis=-1)        	
+    return vol_in, vol_out
+
+def _image_vol_normalization(vol):
+    vol = vol/255
+    vol[vol < 0] = 0
+    vol[vol > 1] = 1
+    return vol
+
+def _image_vol_augmentation(vol_in, vol_out, tr_phase):
+    vol_in, vol_out = _shift_first_frame(vol_in, vol_out, tr_phase)
+    vol_in_norm = _image_vol_normalization(vol_in)
+    vol_out_norm = _image_vol_normalization(vol_out)
+    vol_shape = [sh for sh in vol_in.shape]
+    vol_shape.append(1)
+    vol_in, vol_out = np.empty(vol_shape), np.empty(vol_shape)
+    vol_in[:,:,:,0], vol_out[:,:,:,0] = vol_in_norm, vol_out_norm
+    return [vol_in, vol_out]
+
+def _mask_generation(clt_pos_dta, NFClt):
+    vol_mask = np.zeros((128, 128, 50))
+    for frm in range(50):
+        if NFClt:
+            x0, x1 = clt_pos_dta[0,0], clt_pos_dta[0,1]
+            y0, y1 = clt_pos_dta[1,0], clt_pos_dta[1,1]
+        elif clt_pos_dta.shape[-1] == 1:
+            x0, x1 = clt_pos_dta[frm][0][0,0], clt_pos_dta[frm][0][0,1]
+            y0, y1 = clt_pos_dta[frm][0][1,0], clt_pos_dta[frm][0][1,1]
+        else:
+            x0 = np.min([clt_pos_dta[0][0][0,0], clt_pos_dta[0][1][0,0]])
+            x1 = np.max([clt_pos_dta[0][0][0,1], clt_pos_dta[0][1][0,1]])
+            y0 = np.min([clt_pos_dta[0][0][1,0], clt_pos_dta[0][1][1,0]])
+            y1 = np.max([clt_pos_dta[0][0][1,1], clt_pos_dta[0][1][1,1]])
+        # create the masks
+        vol_mask[x0:x1,y0:y1,frm] = 1
+    return vol_mask
+
+class DataGen(tf.keras.utils.Sequence):
+    """ 
+    Generating batches of input cluttered volumes and their corresponding
+    clutter-free output volumes 
+    """
+    def __init__(
+        self,
+        dim:list,
+        in_dir:str,
+        out_dir:str,
+        id_list:list,
+        batch_size:int,
+        masked_in=True,
+        tr_phase=True,
+	    *args,
+        **kwargs):
+        
+        'Initialization'
+        self.dim = dim
+        self.in_dir = in_dir
+        self.out_dir = out_dir
+        self.id_list = id_list
+        self.batch_size = batch_size
+        self.masked_in = masked_in
+        self.tr_phase = tr_phase
+		
+    def __len__(self):
+        return int(np.floor(len(self.id_list) / self.batch_size))
+    
+    def __getitem__(self, idx):
+        batch = self.id_list[idx*self.batch_size:(idx+1)*self.batch_size]
+        x_aug, y_aug, mask = [], [], []
+        for i, ID in enumerate(batch):
+            vol_in = sio.loadmat(self.in_dir[ID])['data_artf']
+            vol_out = sio.loadmat(self.out_dir[ID])['data_org']           
+            # Call the data augmentation function 
+            aug_vols = _image_vol_augmentation(vol_in, vol_out, self.tr_phase)
+            # Creat the binary mask
+            if self.masked_in:
+                clt_pos_id = self.in_dir[ID]
+                clt_pos_id = clt_pos_id.replace('/data_', '/pos_')
+                if 'pos_NFClt' in clt_pos_id:
+                    clt_pos_dta = sio.loadmat(clt_pos_id)['pos_sp_nf']
+                    NFClt = True
+                else:
+                    clt_pos_dta = sio.loadmat(clt_pos_id)['pos_artf']
+                    NFClt = False
+                vol_mask = _mask_generation(clt_pos_dta, NFClt)
+                mask.append(vol_mask)
+            x_aug.append(aug_vols[0])
+            y_aug.append(aug_vols[1])
+        
+        if self.tr_phase and self.masked_in:
+            return np.asarray(x_aug), np.asarray(y_aug), np.asarray(mask)
+        elif self.tr_phase:
+            return np.asarray(x_aug), np.asarray(y_aug)
+        else:
+            return np.asarray(x_aug)

Filter3D_Lrec_adv/Discriminator.py

@@ -0,0 +1,70 @@
+""" 
+Functions for training discriminator model (ResNet) of the 3D clutter filtering network.
+"""
+import tensorflow as tf
+from tensorflow.keras import backend as K
+from tensorflow.keras.layers import (Conv3D, MaxPooling3D, Dense, Activation, BatchNormalization, Add,
+                                     Dropout, Concatenate, multiply, Input, Flatten, Lambda) 
+from tensorflow.keras.models import Model
+from tensorflow.keras.optimizers import Adam
+
+def max_fnc(x):
+    return K.max(x, axis=1, keepdims=False)
+
+def ResNetModule(x, n_filters, n_blks, krnl_size=3, stride=1, name=None, **prm):
+    
+    shortcut = Conv3D(2*n_filters, 1, strides=stride, name=f"{name}_conv_shortcut")(x)
+    shortcut = BatchNormalization(name=f"{name}_bn_shortcut")(shortcut)
+    for i in range(n_blks):
+        x = Conv3D(n_filters, krnl_size, padding='same', name=f"{name}_conv{i+1}")(x)
+        x = BatchNormalization(name=f"{name}_bn{i+1}")(x)
+        x = Activation('relu', name=f"{name}_relu{i+1}")(x)
+        x = Dropout(0.05, name=f"{name}_drp{i+1}")(x)
+    
+    x = Conv3D(2*n_filters, 1, name=f"{name}_conv_last")(x)
+    x = BatchNormalization(name=f"{name}_bn_last")(x)
+    x = Add(name=f"{name}_add")([shortcut, x])
+    x = Activation('relu', name=f"{name}_out")(x)
+    return x
+
+def ConvToFC(inp, kernel_size=3, name=None, **prm):
+    inp_shape = K.int_shape(inp)
+    x = Conv3D(filters=prm['ConvToFC'][0], kernel_size=(1,1,inp_shape[3]), 
+               padding='valid', activation='relu', strides=1, name=f'{name}_1')(inp)
+    x = Conv3D(filters=prm['ConvToFC'][1], kernel_size=(inp_shape[1],inp_shape[2],1),
+               padding='valid', activation='relu', strides=1, name=f'{name}_2')(x)
+    if len(prm['ConvToFC']) > 2:
+        for i in range(2, len(prm['ConvToFC'])):
+            x = Conv3D(filters=prm['ConvToFC'][i], kernel_size=1, padding='valid',
+                       activation='relu', strides=1, name=f'{name}_{i+1}')(x)
+    print(x.shape)
+    x = Flatten()(x)
+    return x
+
+def ResNet(inp, n_krn, name=None, **prm):    
+    x = Conv3D(n_krn, kernel_size=prm['kernel_size'], padding='same',
+               data_format="channels_last", strides=1, name=f"{name}_conv0")(inp)
+    x = BatchNormalization(name=f"{name}_bn0")(x)
+    x = Activation('relu', name=f"{name}_relu0")(x)
+    for i in range(len(prm['lvl_blks_config'])):        
+        x = MaxPooling3D(pool_size=prm['pool_size'], strides=prm['strides'], name=f"{name}_pool{i+1}")(x)
+        x = ResNetModule(x, n_filters=(2**i)*n_krn, krnl_size=prm['kernel_size'], stride=1, 
+                         n_blks=prm['lvl_blks_config'][i], name=f"{name}_module{i+1}", **prm)       
+    x = ConvToFC(x, name='CnvToFC', **prm)
+    return x
+
+def DenseLayer(x, name=None, **prm):
+  for i in range(len(prm['dense_layer_spec'])):
+    x = Dense(prm['dense_layer_spec'][i], activation='relu', name=f"{name}_DenseLayer_{i}")(x)
+  return Dense(1, activation='sigmoid', name=f"{name}_DenseLayer_out")(x) 
+
+def discriminator_3D(lr, **prm):
+    inp_3D = Input(prm['input_dim'])
+    conv_out = ResNet(inp_3D, n_krn=prm['n_init_filters'], name=prm['model_name'], **prm) 
+    print(conv_out.shape)
+    dense_out = DenseLayer(x=conv_out, name=prm['model_name'], **prm)
+    conv_model = Model(inputs=inp_3D, outputs=dense_out, name=prm['model_name'])
+    conv_model.summary()
+    opt = Adam(learning_rate=lr)    
+    conv_model.compile(optimizer=opt, loss=prm['loss'], metrics=prm['metrics'])     
+    return conv_model

Filter3D_Lrec_adv/TestClutterFilter3D_GAN.py

@@ -0,0 +1,61 @@
+""" 
+Module for testing the 3D clutter filtering model.
+"""
+import os
+import sys
+import argparse
+import json
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.layers import Input 
+from tensorflow.keras.models import Model
+from tensorflow.keras.optimizers import Adam
+
+code_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.dirname(code_dir))
+from Filter3D_Lrec.utils import *
+from Filter3D_Lrec.Error_analysis import *
+from Filter3D_Lrec.Model_ClutterFilter3D import Unet3D
+from DataGen import DataGen
+from Discriminator import discriminator_3D
+from Train_GAN import train_gan
+
+def data_generation(in_ids_te, out_ids_te, config):
+    DtaGenTe_prm = {
+        'dim': config["generator_prm"]["input_dim"],
+        'in_dir': in_ids_te,
+        'out_dir': out_ids_te,
+        'id_list': np.arange(len(in_ids_te)),
+        'batch_size': config["learning_prm"]["batch_size"],
+        'tr_phase': False} 
+    return DataGen(**DtaGenTe_prm)
+
+def generator_3D(**prm):
+    prm_gen = prm
+    prm_gen["network_prm"] = prm["generator_prm"]
+    main_in = Input(prm_gen["network_prm"]["input_dim"])
+    filter_out = Unet3D(x_in=main_in, name="CF", **prm_gen)
+    model = Model(inputs=main_in, outputs=filter_out, name='generator_3D')
+    model.summary()
+    return model
+
+def main(config):
+    in_ids_te, out_ids_te, te_subject, val_subject = id_preparation(config)
+    te_gen = data_generation(in_ids_te, out_ids_te, config)
+    weight_dir = create_weight_dir(val_subject, te_subject, config) 
+    generator = generator_3D(**config)
+    generator.load_weights(os.path.join(weight_dir, config["weight_name"] + ".hdf5"))
+    results_te = generator.predict_generator(te_gen, verbose=2)
+    df_errors = compute_mae(in_ids_te, results_te)
+    df_errors.to_csv(
+        os.path.join(weight_dir, config["weight_name"] + ".csv"))
+    return None
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", help="path of the config file", default="config.json")
+    args = parser.parse_args()
+    assert os.path.isfile(args.config)
+    with open(args.config, "r") as read_file:
+        config = json.load(read_file)
+    main(config)

Filter3D_Lrec_adv/TrainClutterFilter3D_L2_adv_loss.py

@@ -0,0 +1,80 @@
+""" 
+Module for training the 3D clutter filtering model with adversarial loss.
+"""
+import os
+import sys
+import argparse
+import json
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.layers import Input 
+from tensorflow.keras.models import Model
+from tensorflow.keras.optimizers import Adam
+
+code_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.dirname(code_dir))
+from Filter3D_Lrec.utils import *
+from Filter3D_Lrec.Model_ClutterFilter3D import Unet3D
+from DataGen import DataGen
+from Discriminator import discriminator_3D
+from Train_GAN import train_gan
+
+def data_generation(in_ids_tr, out_ids_tr, config):
+    DtaGenTr_prm = {
+        'dim': config["generator_prm"]["input_dim"],
+        'in_dir': in_ids_tr,
+        'out_dir': out_ids_tr,
+        'id_list': np.arange(len(in_ids_tr)),
+        'batch_size': config["learning_prm"]["batch_size"],
+        'tr_phase': True}
+    return DataGen(**DtaGenTr_prm)
+
+def generator_3D(**prm):
+    prm_gen = prm
+    prm_gen["network_prm"] = prm["generator_prm"]
+    main_in = Input(prm_gen["network_prm"]["input_dim"])
+    filter_out = Unet3D(x_in=main_in, name="CF", **prm_gen)
+    model = Model(inputs=main_in, outputs=filter_out, name='generator_3D')
+    model.summary()
+    return model
+
+def gan_3D(gen_model, disc_model, w_g, w_d, lr):
+    disc_model.trainable = False
+    disc_output = disc_model(gen_model.output)
+    opt = Adam(lr=lr, beta_1=0.5)
+    model = Model(gen_model.input, [gen_model.output, disc_output])
+    model.compile(loss=['mse', 'binary_crossentropy'],
+                  loss_weights=[w_g, w_d],
+                  optimizer=opt)
+    return model
+
+def main(config):
+    in_ids_tr, in_ids_val, out_ids_tr, out_ids_val, val_subject, te_subject = id_preparation(config)
+    weight_dir = create_weight_dir(val_subject, te_subject, config)
+    weight_name = (
+      f'CF3D_GAN_ValTeSbj_{val_subject}_{te_subject}'
+      f'_InSkp{config["generator_prm"]["in_skip"]}'
+      f'_Att{config["generator_prm"]["attention"]}_lr{config["learning_prm"]["lr"]}'
+      f'_wg{config["generator_prm"]["w_g"]}'
+      f'_MaskedIn{config["discriminator_prm"]["masked_in"]}')
+    tr_gen = data_generation(in_ids_tr, out_ids_tr, config)
+    generator = generator_3D(**config)
+    discriminator = discriminator_3D(config["learning_prm"]["lr"],
+                                     **config["discriminator_prm"])
+    gan_model = gan_3D(generator, discriminator, 
+                       config["generator_prm"]["w_g"], config["discriminator_prm"]["w_d"],
+                       config["learning_prm"]["lr"])
+    train_gan(g_model=generator, d_model=discriminator, gan_model=gan_model,
+              n_epochs=config["learning_prm"]["n_epochs"],
+              tr_batches=tr_gen, w_dir=weight_dir, w_name=weight_name,
+              masked_in=config["discriminator_prm"]['masked_in'])
+    return None
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", help="path of the config file", default="config.json")
+    args = parser.parse_args()
+    assert os.path.isfile(args.config)
+    with open(args.config, "r") as read_file:
+        config = json.load(read_file)
+    main(config)

Filter3D_Lrec_adv/Train_GAN.py

@@ -0,0 +1,36 @@
+""" 
+Functions for training the 3D GAN model.
+"""
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras import backend as K 
+from tensorflow.keras.models import Model
+
+def _apply_the_mask(in_g, in_d, gen_out, th=0.1):
+    mask = np.abs(in_g-in_d)
+    mask[mask < th] = 0
+    mask[mask >= th] = 1
+    in_d_masked = np.multiply(in_d, mask)
+    gen_out_masked = np.multiply(gen_out, mask)
+    return gen_out_masked, in_d_masked
+
+def train_gan(g_model, d_model, gan_model, n_epochs, tr_batches, w_dir, w_name, masked_in, **prm):
+    for i in range(n_epochs):
+        print(f"epoch:{i}")       
+        rnd_ids = np.random.permutation(tr_batches.__len__())
+        for j in range(tr_batches.__len__()):
+            #  Train Discriminator
+            in_g, in_d = tr_batches.__getitem__(rnd_ids[j])[0], tr_batches.__getitem__(rnd_ids[j])[1]
+            gen_out = g_model.predict(in_g) 
+            if masked_in:
+                gen_out, in_d = _apply_the_mask(in_g, in_d, gen_out)
+            d_loss_r, d_acc_r = d_model.train_on_batch(in_d, np.ones((1, 1)))
+            d_loss_f, d_acc_f = d_model.train_on_batch(gen_out, np.zeros((1, 1)))
+            d_loss = 0.5 * np.add(d_loss_r, d_loss_f)
+            #  Train Generator
+            g_loss = gan_model.train_on_batch(tr_batches.__getitem__(rnd_ids[j])[0],
+                                              [tr_batches.__getitem__(rnd_ids[j])[1], np.ones((1, 1))])
+        # Save weights after each epoch
+        filename = (w_dir + '/' + f"{w_name}_epc{i}_g_loss{np.round(g_loss, 3)}_d_loss{np.round(d_loss, 3)}" + ".hdf5")
+        g_model.save_weights(filename)
+    return None

Filter3D_Lrec_adv/config.json

@@ -0,0 +1,50 @@
+{
+    "paths": {
+        "data_path": "",
+        "save_path": ""
+    },
+    "subject_list": ["rca", "ladprox", "laddist", "lcx", "normal"],
+    "CV": {
+        "val_subject_id": 0
+    },
+    "generator_prm": {
+        "input_dim": [128, 128, 50, 1],
+        "n_levels": 4,
+        "n_init_filters": 16,
+        "in_skip": true,
+        "attention": true,
+        "kernel_size": 3,
+        "conv_stride": 1,
+        "upsampling_stride": [2, 2, 1],
+        "pool_size": [2, 2, 1], 
+        "pool_stride": 1,
+        "padding": "same",
+        "act": "linear",
+        "w_g": 0.999
+    },
+    "discriminator_prm": {
+        "input_dim": [128, 128, 50, 1],
+        "n_init_filters": 16,
+        "lvl_blks_config": [3, 4, 6, 3],
+        "kernel_size": 3,
+        "pool_size": 2,
+        "strides": 2,
+        "ConvToFC": [128, 32],
+        "dense_layer_spec": [64, 32, 16],
+        "loss": "binary_crossentropy",
+        "metrics": ["accuracy"],
+        "attention": false,
+        "masked_in": true,
+        "model_name": "disc_ResNet34",
+        "w_d": 0.001
+    },
+    "learning_prm": {
+        "batch_size": 1,
+        "lr": 1e-4,
+        "drp": 0.05,
+        "n_epochs": 10
+    },
+    "tr_phase": true,
+    "model_name": "CF3D_AdvLoss",
+    "weight_name": ""
+}