main_testing_VCT.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Feb 11 14:20:29 2022

@author: Rodrigo
"""

import matplotlib.pyplot as plt
import numpy as np
import torch
import pydicom
import os
import pathlib

from tqdm import tqdm

# Own codes
from libs.models import Gen
from libs.utilities import load_model_gan, makedir
from libs.dataset import scale, de_scale

#%%
def img2rois(img_ld):
    
    h, w = img_ld.shape
    
    # How many Rois fit in the image?
    n_h = h % 64
    n_w = w % 64
    
    if n_h == 0:
        h_pad = h
    else:
        h_pad = (h//64 + 1) * 64
        
    if n_w == 0:
        w_pad = w
    else:
        w_pad = (w//64 + 1) * 64
    
    # Calculate how much padding is necessary and sum 64 for the frontiers
    padding = (((h_pad - h)//2 + 64, (h_pad - h)//2 + 64),
               ((w_pad - w)//2 + 64, (w_pad - w)//2 + 64))
    
    # Pad the image
    img_ld_pad = np.pad(img_ld, padding, mode='reflect')
            
    n_h = h_pad // 64 
    n_w = w_pad // 64 
    
    # Allocate memory to speed up the for loop
    rois = np.empty((n_h*n_w, 1, 192, 192), dtype='float32')

    nRoi = 0
    # Get the ROIs
    for i in range(n_h):
        for j in range(n_w):
            rois[nRoi, 0, :, :] = img_ld_pad[i*64: (i+3)*64, j*64:(j+3)*64]
            nRoi += 1
            
    return rois, img_ld_pad.shape
            
def rois2img(rst_rois, original_shape, padded_shape):

    
    rst_img = np.empty((padded_shape))
    
    n_h = (padded_shape[0] // 64) - 2
    n_w = (padded_shape[1] // 64) - 2
      
    nRoi = 0
    # Reconstruct image format
    for i in range(n_h):
        for j in range(n_w):
            rst_img[(i+1)*64:(i+2)*64, (j+1)*64:(j+2)*64] = rst_rois[nRoi,0,64:128,64:128]
            nRoi += 1
        
    org_h, org_w = original_shape
    pad_h, pad_w = padded_shape
    
    # How much to crop?
    start_w = (pad_w - org_w) // 2
    start_h = (pad_h - org_h) // 2
    
    # Crop image
    rst_img = rst_img[start_h:start_h+org_h, start_w:start_w+org_w]
        
    return rst_img
         
def model_forward(model, img_ld, red_factor, batch_size, vmin, vmax):
    
    # Change model to eval
    model.eval()

    # Normalize image
    img_ld = scale(img_ld, red_factor=red_factor, vmin=vmin, vmax=vmax)    

    # Extract ROIs
    rois, padded_shape = img2rois(img_ld)
    
    # Allocate memory to speed up the for loop 
    rst_rois = np.empty_like(rois)
    
    for x in range(0,rois.shape[0],batch_size):
        
        # Get the batch and send to GPU
        batch = torch.from_numpy(rois[x:x+batch_size]).to(device)
        
        # Forward through the model
        with torch.no_grad():
            batch = model(batch)
        
        # Get from GPU
        rst_rois[x:x+batch_size] = batch.to('cpu').numpy()

    # Contruct the image
    rst_img = rois2img(rst_rois, img_ld.shape, padded_shape)
    
    # Normalize image (Inv)
    rst_img = de_scale(rst_img, vmin=vmin, vmax=vmax)
    
    return rst_img

def test(model, data2read, path2write, red_factor, mAsLowDose, batch_size, nrlz, vmin, vmax):
    
    
    for folder_name in tqdm(data2read):
        
        patient_name = folder_name.split('/')[-1]
        
        for rlz in range(1,nrlz+1):
    
            folder_name_wc = '/'.join(folder_name.split('/')[:-1]) + '/noisy/' + patient_name + '-{}mAs-rlz{}-wc'.format(90,rlz)
    
            file_names = list(pathlib.Path(folder_name_wc).glob('**/*'))
        
            for file_name in file_names:
                
                file_name = str(file_name) 
                
                # Read dicom image
                dcmH = pydicom.dcmread(file_name)
        
                # Read dicom image pixels
                img_ld = dcmH.pixel_array.astype('float32')
             
                # Forward through model
                rst_img = model_forward(model, img_ld, red_factor, batch_size, vmin, vmax)
        
                folder2write_name = path2write + patient_name + '-{}mAs-rlz{}-uc-{}'.format(mAsLowDose,rlz,target_mAs_arg_fine) 
                file2write_name = file_name.split('/') [-1]
                
                # Create output dir (if needed)
                makedir(folder2write_name)
                
                # Copy the restored data to the original dicom header
                dcmH.PixelData = np.uint16(rst_img)
                
                # Write dicom
                pydicom.dcmwrite(os.path.join(folder2write_name,file2write_name),
                                 dcmH, 
                                 write_like_original=True) 

    return

#%%

if __name__ == '__main__':
      
    # Noise scale factor
    mAsFullDose = 360
    mAsLowDose = 90
    
    nrlz = 10
    
    batch_size = 50
    
    
    if mAsFullDose == 360:
        target_mAs_arg = '_360mAs'
        target_mAs_arg_fine = 'to360mAs'
        # 360mAs
        vmin=2000
        vmax=18000
    else:
        target_mAs_arg = ''
        target_mAs_arg_fine = ''
        # 90mAs
        vmin=0
        vmax=7000
        
    
    red_factor = mAsLowDose / mAsFullDose
    
    path_models = "final_models/"
    path2write = "data/outputs/"
    
    data2read = ['/media/rodrigo/Data/images/UPenn/Phantom/VCT/VCT_Bruno_500/GE-projs/Jhones_888076.1.555465657049.20200828021423277-proj']
    
    makedir(path2write)
    
    path_final_model = path_models + 'cGAN_ResNet_Decor-90mAs{}.pth'.format(target_mAs_arg)
    
    # Test if there is a GPU
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    
    # Create models
    model = Gen() 
    
    # Send it to device (GPU if exist)
    model = model.to(device)
    
    # Load gen pre-trained model parameters (if exist)
    _ = load_model_gan([model, None, None, None],path_final_model=path_final_model)
    
    print("Running test on {}.".format(device))
    
    test(model, data2read, path2write, red_factor, mAsLowDose, batch_size, nrlz, vmin, vmax)