DIVA_Cardiac_USimage_Enhancement.ipynb

{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"provenance":[],"mount_file_id":"14eUzPa-tZMOCIgY6DpufDYcnOca5GGvh","authorship_tag":"ABX9TyMH96Gs3KY2Va7FOoA8C7vD"},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"}},"cells":[{"cell_type":"code","execution_count":null,"metadata":{"id":"VXGA8GOHkmE0"},"outputs":[],"source":["import argparse\n","from argparse import ArgumentParser\n","import glob\n","import cv2\n","import re\n","import os, glob, datetime\n","import numpy as np\n","import tensorflow as tf\n","import numpy as np\n","import time\n","import math\n","from keras.layers import  Input,Conv2D,BatchNormalization,Activation,Subtract, Reshape\n","from keras.models import Model, load_model\n","from tensorflow.python.keras.utils import conv_utils\n","from keras.callbacks import CSVLogger, ModelCheckpoint, LearningRateScheduler\n","from tensorflow.keras.optimizers import Adam\n","import keras.backend as K\n","\n","import skimage\n","from skimage.metrics import structural_similarity as ssim\n","from skimage.io import imread, imsave\n","\n","\n","##--------------------------------------------------------------------------------------------------------\n","##--------------------------------------------------------------------------------------------------------\n","\n","\n","#########--------------           IMPORTANT NOTE       ----------------###############\n","\n","# Please download the pretrained model from \n","#   https://www.dropbox.com/s/6edzcf76lnfn7uj/DIVA_US_image_despeckling_model.hdf5?dl=0\n","#   and save file into 'Model' folder before testing\n","\n","# Select the path of the directory properly to give access to the models, datasets and results\n","directory_path =  '___________________________________' #'Choose your directory\n","\n","\n","\n","\n","parser = argparse.ArgumentParser()\n","# choose if needed\n","parser.add_argument('--model', default='DIVA', type=str, help='choose a type of model')\n","\n","parser.add_argument('--set_dir', default=os.path.join(directory_path ,'Data/'), type=str, help='directory of test dataset')\n","parser.add_argument('--Speckled_set_names', default=['test_Images'], type=list, help='name of the speckled test dataset')\n","\n","\n","#########--------------           IMPORTANT NOTE       ----------------###############\n","\n","# read the clean data set if available\n","#parser.add_argument('--set_names_clean', default='More_test_Images_Original', type=str, help='name of the clean test dataset')\n","\n","parser.add_argument('--model_dir', default=os.path.join(directory_path ,'Model'), type=str, help='directory of the model')\n","                                                     \n","parser.add_argument('--model_name', default='DIVA_US_image_despeckling_model.hdf5', type=str, help='the model name')\n","parser.add_argument('--result_dir', default=os.path.join(directory_path,'results'), type=str, help='directory of results')\n","parser.add_argument('--save_result', default=1, type=int, help='save the denoised image, 1 or 0')\n","\n","parser.add_argument('-f', '--file', required=False)\n","args = parser.parse_args()\n","\n","\n","\n","print(args.model_name)\n","\n","## -------------------------------------------------------------------------------------------------------------\n","\n","##------------------------------------------------------------------------------------------------------------\n","\n","\n","# Define useful functions\n","\n","def to_tensor(img):\n","    if img.ndim == 2:\n","        return img[np.newaxis,...,np.newaxis]\n","    elif img.ndim == 3:\n","        return np.moveaxis(img,2,0)[...,np.newaxis]\n","\n","def from_tensor(img):\n","    return np.squeeze(np.moveaxis(img[...,0],0,-1))\n","\n","def log(*args,**kwargs):\n","     print(datetime.datetime.now().strftime(\"%Y-%m-%d %H:%M:%S:\"),*args,**kwargs)\n","\n","def save_result(result,path):\n","    path = path if path.find('.') != -1 else path+'.png'\n","    ext = os.path.splitext(path)[-1]\n","    if ext in ('.txt','.dlm'):\n","        np.savetxt(path,result,fmt='%2.4f')\n","    else:\n","        imsave(path,np.clip(result,0,1))\n","\n","\n","def show(x,title=None,cbar=False,figsize=None):\n","    import matplotlib.pyplot as plt\n","    plt.figure(figsize=figsize)\n","    plt.imshow(x,interpolation='nearest',cmap='gray')\n","    if title:\n","        plt.title(title)\n","    if cbar:\n","        plt.colorbar()\n","    plt.show()\n","\n","def psnr(target, ref):\n","    # Assume target is RGB/BGR image\n","    target_data = target.astype(np.float32)\n","    ref_data = ref.astype(np.float32)\n","    \n","    diff = ref_data - target_data\n","    diff = diff.flatten('C')\n","    \n","    rmse = np.sqrt(np.mean(diff ** 2.))\n","    \n","    return 20 * np.log10(1. / rmse)\n","\n","\n","def snr(target, ref):\n","    # Assume target is RGB/BGR image\n","    target_data = target.astype(np.float32)\n","    ref_data = ref.astype(np.float32)\n","    \n","    diff = ref_data - target_data\n","    diff = diff.flatten('C')\n","    target_data = target_data.flatten('C')\n","    \n","    rmse_diff = np.sqrt(np.mean(diff ** 2.))\n","    rmse_target = np.sqrt(np.mean(target_data ** 2.))\n","\n","    return 20 * np.log10(rmse_target / rmse_diff)\n","\n","##----------------------------------------------------------------------------------------------------------------\n","\n","##----------------------------------------------------------------------------------------------------------------\n","\n","\n","if __name__ == '__main__':    \n","    \n","    # Load trained model\n","    print('Model- ',args.model)\n","    model = load_model(os.path.join(args.model_dir, args.model_name),compile=False)\n","    log('load trained model')\n","    \n","    # Make directory for results\n","    if not os.path.exists(args.result_dir):\n","        os.mkdir(args.result_dir)\n","\n","#-----------------------------------------------------------------------------------------------------------------------------------\n","#-----------------------------------------------------------------------------------------------------------------------------------\n","\n","\n","    for set_cur in args.Speckled_set_names:  \n","        \n","        if not os.path.exists(os.path.join(args.result_dir,set_cur)):\n","            os.mkdir(os.path.join(args.result_dir,set_cur))\n","\n","        \n","        for im in os.listdir(os.path.join(args.set_dir,set_cur)): \n","            if im.endswith(\".jpg\") or im.endswith(\".bmp\") or im.endswith(\".png\"):\n","\n","                # speckled noise case\n","                noise_type = 'Speckled'\n","\n","                # read the speckled image\n","                y = np.array(imread(os.path.join(args.set_dir,set_cur,im)), dtype=np.float32) / 255.0\n","\n","\n","                y = y.astype(np.float32)\n","                y_  = to_tensor(y)\n","                start_time = time.time()\n","                x_ = model.predict(y_)      # Despeckling process\n","                elapsed_time = time.time() - start_time\n","                print('%10s : %10s : %2.4f second'%(set_cur,im,elapsed_time))\n","\n","                # Despeckled image\n","                x_=from_tensor(x_)\n","\n","\n","                if args.save_result:\n","                    name, ext = os.path.splitext(im)\n","                    show(np.hstack((y,x_))) # show the image\n","                    save_result(x_,path=os.path.join(args.result_dir,set_cur,name+ext)) # save the Despeckled image\n","                \n","#-----------------------------------------------------------------------------------------------------------------------------------\n","#-----------------------------------------------------------------------------------------------------------------------------------\n","\n"]}]}