inference_args.py

import torch
from utils.image_processing import change_orientation
import numpy as np
from utils.nifti2npy import Nifti2Npy
from scipy.ndimage import gaussian_filter
import albumentations as A
from model.layers import BPRNetV1
import nibabel as nib
import torchvision
import matplotlib.pyplot as plt
import json
from scipy.interpolate import interp1d
from torch.nn import functional as F
from utils.chart import ChartGroup, Chart
from utils.nifty2dicom import dicom_to_nifti
import os
import argparse

parser = argparse.ArgumentParser(description="Process input arguments.")
parser.add_argument("--input_path", type=str, help="Input nifti file path")
parser.add_argument("--output_path", type=str, help="Result saving path")
parser.add_argument("--case_name", type=str, help="Case name used for saving the result", default='')
parser.add_argument("--device", type=str, help="Device to run", default='cuda')

args = parser.parse_args()

def plot_soft_slice(x1x2, chart, weight=None):

    if weight is None:
        weight = np.ones(x1x2.shape[0])

    r = 1
    for i in range(4): 
        theta = np.linspace(-np.pi / 4 + i * (np.pi / 2), np.pi / 4 + i * (np.pi / 2), 40)
        circle_x = r * np.sin(theta)
        circle_y = r * np.cos(theta)
        chart.line(circle_x, circle_y, color=['red', 'blue', 'orange', 'green'][i])
    chart.line([0, 0], [-1, 1], color='black')
    chart.line([-1, 1], [0, 0], color='black')
    chart.ax.text(-0.2, 0.85, 'Non-Con')
    chart.ax.text(0.5, 0.1, 'Arterial')
    chart.ax.text(-0.9, 0.1, 'Delayed')
    chart.ax.text(-0.3, -0.85, 'Portalvenous')

    for coords, w in zip(x1x2, weight):
        coords /= np.linalg.norm(coords)
        chart.scatter([coords[0]], [coords[1]], color='black', markersize=6 * w, opacity=0.3)

def build_bpr_model(checkpoint):
    model = BPRNetV1()
    model.load_state_dict(torch.load(checkpoint)['model'])
    return model

def build_epr_model(checkpoint):
    model_bpr = BPRNetV1()
    D = 1024
    model = torchvision.models.vgg16()
    model.classifier = torch.nn.Sequential(
        torch.nn.Linear(25088, D, bias=True),
        torch.nn.ReLU(inplace=True),
        torch.nn.Dropout(p=0.5, inplace=False),
        torch.nn.Linear(D, D, bias=True),
        torch.nn.ReLU(inplace=True),
        torch.nn.Dropout(p=0.5, inplace=False),
        torch.nn.Linear(D, 2, bias=True)
    )
    model.features = model_bpr.slice_encoder
    model.load_state_dict(torch.load(checkpoint)['model'])
    return model

def predict(case_name, nifti_path, bpr_net, epr_net, voting_weight, device='cuda'):

    bpr_net.eval()
    epr_net.eval()

    # preprocess nifti

    if nifti_path.endswith('.nii') or nifti_path.endswith('.nii.gz'):
        nifti = nib.load(nifti_path)
    else:
        # dicom
        nifti = dicom_to_nifti(nifti_path)
    
    nifti_LAS = change_orientation(nifti, orientation_code=['L', 'A', 'S'])

    vol = np.array(nifti_LAS.dataobj).astype(np.float32)

    p = Nifti2Npy()
    vol = p.rescale_xy(vol)
    vol = 2 * (vol - vol.min()) / (vol.max() - vol.min()) - 1
    resize = A.Resize(128, 128)
    blurred = gaussian_filter(vol, (0.8, 0.8, 0), truncate=3)
    slices_128norm = resize(image=blurred)['image']
    slices_128norm = slices_128norm.astype(np.float32)

    slices = np.transpose(slices_128norm, [2, 0, 1])
    slices = torch.from_numpy(slices).float().to(device)
    
    # run bpr to get position scores
    x = slices[None]
    batch_size = x.shape[0]
    num_slices = x.shape[1]

    N = 200
    seg_start = 0
    seg_end = N

    scores_all_segments = []
    x = x.reshape(batch_size * num_slices, 1, x.shape[2], x.shape[3])
    while seg_start < num_slices:
        seg_end = min(seg_end, num_slices)
        x_seg = x[seg_start: seg_end, :]
        with torch.no_grad():
            scores_seg = bpr_net(x_seg)
        scores_all_segments.append(scores_seg)
        seg_start += N
        seg_end += N

    scores = torch.cat(scores_all_segments, dim=0)
    scores = scores.detach().cpu().numpy()[:, 0]

    # run epr to get phase direction (x1 x2)
    x = slices
    x = x[:, None, :, :]
    x1x2_all_segments = []
    num_slices = x.shape[0]
    N = 200
    seg_start = 0
    seg_end = N

    while seg_start < num_slices:
        seg_end = min(seg_end, num_slices)
        x_seg = x[seg_start: seg_end]
        with torch.no_grad():
            x_seg = epr_net.features(x_seg)
            x_seg = epr_net.avgpool(x_seg)
            x_seg = torch.flatten(x_seg, 1)
            for i in range(6):
                x_seg = epr_net.classifier[i](x_seg)
            x1x2_seg = epr_net.classifier[6](x_seg)
        x1x2_all_segments.append(x1x2_seg)
        seg_start += N
        seg_end += N

    x1x2 = torch.cat(x1x2_all_segments, dim=0).detach().cpu().numpy()

    # run positonal-weighted voting to get phase label
    votes = []
    slice_weights = {}
    weighted_pred = {}
    for phase in [0, 1, 2, 3]:
        cur_phase_weight = voting_weight[str(phase)]
        interp = interp1d(cur_phase_weight['score'], cur_phase_weight['weight'], fill_value=0, bounds_error=False)
        weight = interp(scores) * 10 # alpha
        weight = F.softmax(torch.from_numpy(weight), dim=0).numpy() * weight.shape[0]
        slice_weights[phase] = weight
        coords = x1x2 / np.linalg.norm(x1x2, axis=1)[:, None]
        weighted_avg_coords = (coords * weight[:, None]).mean(axis=0)
        weighted_avg_coords /= np.linalg.norm(weighted_avg_coords)
        weighted_pred[phase] = weighted_avg_coords
        V = np.array([[0, 1], [np.cos(-np.pi / 8), np.sin(-np.pi / 8)], [0, -1], [-1, 0]]).T
        weighted_cos = weighted_avg_coords @ V
        votes.append(weighted_cos[phase])

    phase_label = np.array(votes).argmax()

    # generate results and plots
    
    if not os.path.exists(f'{args.output_path}/results'):
        os.mkdir(f'{args.output_path}/results')

    if not os.path.exists(f'{args.output_path}/plots'):
        os.mkdir(f'{args.output_path}/plots')

    res = {'scores': scores.tolist(), 'phase_label': int(phase_label), 'case': case_name,
           'weight': (slice_weights[phase_label] / slice_weights[phase_label].max()).tolist(),
           'x1x2': x1x2.tolist(), 'pred_x1x2': weighted_pred[phase_label].tolist()
           }
    json.dump(res, open(f'{args.output_path}/results/{case_name}.json', 'w'))

    cg = ChartGroup(2, 2)
    c11 = cg.get_chart(1, 1).title('Coronal View of Body CT')
    c12 = cg.get_chart(1, 2).title('BPR scores and voting weights')
    c12.xlabel('position score')
    c12.ylabel('axial slice index')
    c21 = cg.get_chart(2, 1).axis_off().title('Slice-level phase prediction')
    c22 = cg.get_chart(2, 2).axis_off().title('Weighted and final phase prediction')

    voxsz = [abs(nifti_LAS.affine[0, 0]), abs(nifti_LAS.affine[1, 1]), abs(nifti_LAS.affine[2, 2])]
    vol = np.array(nifti_LAS.dataobj)
    vol[vol < -125] = -125
    vol[vol > 250] = 250
    c11.slice(vol, orientation='y', voxsz=voxsz)
    c12.scatter(scores, np.array(range(scores.shape[0])), color=0)
    c12.line(slice_weights[phase_label] / slice_weights[phase_label].max(), np.array(range(scores.shape[0])), color=1)
    plot_soft_slice(x1x2, chart=c21, weight=None)
    plot_soft_slice(x1x2, chart=c22, weight=slice_weights[phase_label] / slice_weights[phase_label].max())

    c22.line([0, weighted_pred[phase_label][0]], [0, weighted_pred[phase_label][1]], color='black')

    plt.savefig(f'{args.output_path}/plots/{case_name}.png')


file_dir = '/'.join(__file__.split('/')[: -1])

bpr_net = build_bpr_model(f'{file_dir}/checkpoints/bpr_model.pth')
epr_net = build_epr_model(f'{file_dir}/checkpoints/epr_model.pth')

voting_weight = json.load(open(f'{file_dir}/epc_voting_weight.json'))

bpr_net.to(args.device)
epr_net.to(args.device)

if args.case_name == '':
    case_name = args.input_path.split('/')[-1]
else:
    case_name = args.case_name

predict(case_name, args.input_path, bpr_net, epr_net, voting_weight, args.device)