示例#1
0
def test():
    model = ModelFlow_stride(2, 3, opt.start_channel).cuda()
    transform = SpatialTransform().cuda()

    model.load_state_dict(torch.load(opt.modelpath))
    model.eval()
    transform.eval()

    grid = generate_grid(imgshape)
    grid = Variable(torch.from_numpy(np.reshape(grid, (1, ) +
                                                grid.shape))).cuda().float()

    start = timeit.default_timer()

    A = Variable(torch.from_numpy(load_5D(opt.fixed))).cuda().float()
    B = Variable(torch.from_numpy(load_5D(opt.moving))).cuda().float()
    start2 = timeit.default_timer()
    print('Time for loading data: ', start2 - start)
    pred = model(A, B)
    F_AB = pred.permute(0, 2, 3, 4, 1).data.cpu().numpy()[0, :, :, :, :]
    F_AB = F_AB.astype(np.float32) * range_flow
    warped_A = transform(A,
                         pred.permute(0, 2, 3, 4, 1) * range_flow,
                         grid).data.cpu().numpy()[0, 0, :, :, :]
    start3 = timeit.default_timer()
    print('Time for registration: ', start3 - start2)

    warped_F_BA = transform(-pred,
                            pred.permute(0, 2, 3, 4, 1) * range_flow,
                            grid).permute(0, 2, 3, 4,
                                          1).data.cpu().numpy()[0, :, :, :, :]
    warped_F_BA = warped_F_BA.astype(np.float32) * range_flow

    start4 = timeit.default_timer()
    print('Time for generating inverse flow: ', start4 - start3)

    save_flow(F_AB, savepath + '/flow_A_B.nii.gz')
    save_flow(warped_F_BA, savepath + '/inverse_flow_B_A.nii.gz')
    save_img(warped_A, savepath + '/warped_A.nii.gz')

    start5 = timeit.default_timer()
    print('Time for saving results: ', start5 - start4)
    del pred

    pred = model(B, A)
    F_BA = pred.permute(0, 2, 3, 4, 1).data.cpu().numpy()[0, :, :, :, :]
    F_BA = F_BA.astype(np.float32) * range_flow
    warped_B = transform(B,
                         pred.permute(0, 2, 3, 4, 1) * range_flow,
                         grid).data.cpu().numpy()[0, 0, :, :, :]
    warped_F_AB = transform(-pred,
                            pred.permute(0, 2, 3, 4, 1) * range_flow,
                            grid).permute(0, 2, 3, 4,
                                          1).data.cpu().numpy()[0, :, :, :, :]
    warped_F_AB = warped_F_AB.astype(np.float32) * range_flow
    save_flow(F_BA, savepath + '/flow_B_A.nii.gz')
    save_flow(warped_F_AB, savepath + '/inverse_flow_A_B.nii.gz')
    save_img(warped_B, savepath + '/warped_B.nii.gz')
def test():
    model = SYMNet(2, 3, opt.start_channel).cuda()
    transform = SpatialTransform().cuda()

    diff_transform = DiffeomorphicTransform(time_step=7).cuda()
    com_transform = CompositionTransform().cuda()

    model.load_state_dict(torch.load(opt.modelpath))
    model.eval()
    transform.eval()
    diff_transform.eval()
    com_transform.eval()

    grid = generate_grid(imgshape)
    grid = torch.from_numpy(np.reshape(grid,
                                       (1, ) + grid.shape)).cuda().float()

    use_cuda = True
    device = torch.device("cuda" if use_cuda else "cpu")

    fixed_img = load_4D(fixed_path)
    moved_img = load_4D(moving_path)

    fixed_img = torch.from_numpy(fixed_img).float().to(device).unsqueeze(dim=0)
    moved_img = torch.from_numpy(moved_img).float().to(device).unsqueeze(dim=0)

    with torch.no_grad():
        F_xy, F_yx = model(fixed_img, moved_img)

        F_X_Y_half = diff_transform(F_xy, grid, range_flow)
        F_Y_X_half = diff_transform(F_yx, grid, range_flow)

        F_X_Y_half_inv = diff_transform(-F_xy, grid, range_flow)
        F_Y_X_half_inv = diff_transform(-F_yx, grid, range_flow)

        F_X_Y = com_transform(F_X_Y_half, F_Y_X_half_inv, grid, range_flow)
        F_Y_X = com_transform(F_Y_X_half, F_X_Y_half_inv, grid, range_flow)

        F_BA = F_Y_X.permute(0, 2, 3, 4, 1).data.cpu().numpy()[0, :, :, :, :]
        F_BA = F_BA.astype(np.float32) * range_flow

        F_AB = F_X_Y.permute(0, 2, 3, 4, 1).data.cpu().numpy()[0, :, :, :, :]
        F_AB = F_AB.astype(np.float32) * range_flow

        warped_B = transform(moved_img,
                             F_Y_X.permute(0, 2, 3, 4, 1) * range_flow,
                             grid).data.cpu().numpy()[0, 0, :, :, :]
        warped_A = transform(fixed_img,
                             F_X_Y.permute(0, 2, 3, 4, 1) * range_flow,
                             grid).data.cpu().numpy()[0, 0, :, :, :]

        save_flow(F_BA, savepath + '/wrapped_flow_B_to_A.nii.gz')
        save_flow(F_AB, savepath + '/wrapped_flow_A_to_B.nii.gz')
        save_img(warped_B, savepath + '/wrapped_norm_B_to_A.nii.gz')
        save_img(warped_A, savepath + '/wrapped_norm_A_to_B.nii.gz')

        print("Finished.")
示例#3
0
def test():
    imgshape_4 = (160 / 4, 192 / 4, 144 / 4)
    imgshape_2 = (160 / 2, 192 / 2, 144 / 2)

    model_lvl1 = Miccai2020_LDR_laplacian_unit_disp_add_lvl1(2, 3, start_channel, is_train=True, imgshape=imgshape_4,
                                               range_flow=range_flow).cuda()
    model_lvl2 = Miccai2020_LDR_laplacian_unit_disp_add_lvl2(2, 3, start_channel, is_train=True, imgshape=imgshape_2,
                                          range_flow=range_flow, model_lvl1=model_lvl1).cuda()

    model = Miccai2020_LDR_laplacian_unit_disp_add_lvl3(2, 3, start_channel, is_train=False, imgshape=imgshape,
                                          range_flow=range_flow, model_lvl2=model_lvl2).cuda()

    transform = SpatialTransform_unit().cuda()

    model.load_state_dict(torch.load(opt.modelpath))
    model.eval()
    transform.eval()

    grid = generate_grid_unit(imgshape)
    grid = torch.from_numpy(np.reshape(grid, (1,) + grid.shape)).cuda().float()

    use_cuda = True
    device = torch.device("cuda" if use_cuda else "cpu")

    fixed_img = load_4D(fixed_path)
    moving_img = load_4D(moving_path)

    fixed_img = torch.from_numpy(fixed_img).float().to(device).unsqueeze(dim=0)
    moving_img = torch.from_numpy(moving_img).float().to(device).unsqueeze(dim=0)

    with torch.no_grad():
        F_X_Y = model(moving_img, fixed_img)

        X_Y = transform(moving_img, F_X_Y.permute(0, 2, 3, 4, 1), grid).data.cpu().numpy()[0, 0, :, :, :]

        F_X_Y_cpu = F_X_Y.data.cpu().numpy()[0, :, :, :, :].transpose(1, 2, 3, 0)
        F_X_Y_cpu = transform_unit_flow_to_flow(F_X_Y_cpu)

        save_flow(F_X_Y_cpu, savepath+'/warpped_flow.nii.gz')
        save_img(X_Y, savepath+'/warpped_moving.nii.gz')

    print("Finished")
示例#4
0
def testOnePair(fixed_path,moving_path,saveResult=1,inputs=[],isSeg=1):
    #support nii, nii.gz, and nrrd
    print(fixed_path)
    print(moving_path)

    ext = ".nii.gz" if ".nii.gz" in fixed_path else (".nii" if ".nii" in fixed_path else ".nrrd" )

    fixedName  = fixed_path.split('/')[-1][:-len(ext)]
    movingName = moving_path.split('/')[-1][:-len(ext)]

    if len(inputs)>0:
        print("input test images are loaded................")
        moving_img = inputs[0]
        fixed_img  = inputs[1]
        fixed_img  = fixed_img.float().to(device)
        moving_img = moving_img.float().to(device)
    else:
       print("loading input test images ................")
       fixed_img  = load_4D(fixed_path)
       moving_img = load_4D(moving_path)
       fixed_img  = torch.from_numpy(fixed_img).float().to(device).unsqueeze(dim=0)
       moving_img = torch.from_numpy(moving_img).float().to(device).unsqueeze(dim=0)


    if isSeg:
        fixed_img[fixed_img>0]=1.0
        moving_img[moving_img > 0] = 1.0

    with torch.no_grad():
        #get displacement field

        displacement_field_tensor = model(moving_img, fixed_img)
        # convert to array for saving to file later
        displacement_field_tensor_cpu = displacement_field_tensor.data.cpu().numpy()[0, :, :, :, :].transpose(1, 2, 3, 0)
        displacement_field_array      = transform_unit_flow_to_flow(displacement_field_tensor_cpu)

        # get transformed image
        transformed_moving_array = transform(moving_img, displacement_field_tensor.permute(0, 2, 3, 4, 1), grid).data.cpu().numpy()[0, 0, :, :, :]
        if not isSeg:
           transformed_moving_array =  normalizeAB(transformed_moving_array,a=-500,b=800)

        print( "len unique transformed img : ", len(np.unique(transformed_moving_array) )  )
        print( "min max transformed img : ", np.min(transformed_moving_array) ,np.max(transformed_moving_array)    )

        print("transformed_moving_array     : ", transformed_moving_array.shape)
        if not isSeg:
            # if segmentation found, transform segmentation as welll
            seg_path = moving_path[:-len(ext)]+'_seg'+ext
            if os.path.isfile(seg_path) and saveResult:
               seg_img =  load_4D(seg_path)
               seg_img = torch.from_numpy(seg_img).float().to(device).unsqueeze(dim=0)
               transformedSeg      = transform(seg_img, displacement_field_tensor.permute(0, 2, 3, 4, 1), grid).data.cpu().numpy()[0, 0, :, :, :]
               print("len unique transformedSeg img : ", len(np.unique(transformedSeg)))
               print("min max transformedSeg img    : ", np.min(transformedSeg), np.max(transformedSeg))
               outputMovingSegPath = savepath + '/' + movingName + '_seg_warpped_moving'+ext
               save_img (transformedSeg              , outputMovingSegPath)

        # save result
        if saveResult:
            outputFixedPath     = savepath +'/'+ fixedName + ext
            outputDispFieldPath = savepath +'/'+ movingName  + '_warpped_flow'  + ext
            outputMovingPath    = savepath +'/'+ movingName  + '_warpped_moving'+ ext
            print("outputFixedPath     : " ,outputFixedPath )
            print("outputDispFieldPath : " ,outputDispFieldPath )
            print("outputMovingPath    : " ,outputMovingPath )

            shutil.copyfile(fixed_path , outputFixedPath)
            #save_flow(displacement_field_array    , outputDispFieldPath)
            save_img_3d(displacement_field_array, fixed_path,outputDispFieldPath)

            #save_img (transformed_moving_array    , outputMovingPath)
            save_img_3d(transformed_moving_array, fixed_path,outputMovingPath)
        #evaluation:
        # compute dice
        fixedSeg = 0.0
        transformedMovingSeg = 0.0
        dicePair = 0.0

        #dicePair = diceMetric(fixedSeg,transformedMovingSeg)
        print(ok)
        return dicePair
    print("Finished")
示例#5
0
def test():
    device = torch.device("cuda:6")
    model = ModelFlow_stride(2, 3, opt.start_channel).cuda(device)
    # model =ModelFlow_stride(2,3,opt.start_channel).cpu()
    # transform = SpatialTransform().cpu()

    transform = SpatialTransform().cuda(device)
    # model.load_state_dict(torch.load(opt.modelpath))
    # model.load_state_dict(torch.load(opt.modelpath, map_location=torch.device('cpu')))
    model.load_state_dict(torch.load(opt.modelpath, map_location=device))

    model.eval()
    transform.eval()
    grid = generate_grid(imgshape)
    grid = Variable(torch.from_numpy(np.reshape(
        grid, (1, ) + grid.shape))).cuda(device).float()
    # grid = Variable(torch.from_numpy(np.reshape(grid, (1,) + grid.shape))).cpu().float()

    start = timeit.default_timer()
    A = sitk.GetArrayFromImage(sitk.ReadImage(opt.fixed, sitk.sitkFloat32))
    B = sitk.GetArrayFromImage(sitk.ReadImage(opt.moving, sitk.sitkFloat32))

    #A = Variable(torch.from_numpy(A)).cuda(device).float()
    #B = Variable(torch.from_numpy(B)).cuda(device).float()

    # A, B = padding(A,B)
    #A = load_5D(A).cuda(device).float()
    #B = load_5D(B).cuda(device).float()

    A = load_5D(A)
    B = load_5D(B)

    A = Variable(torch.from_numpy(A)).cuda(device).float()
    B = Variable(torch.from_numpy(B)).cuda(device).float()

    #A = Variable(torch.from_numpy( load_5D(opt.fixed))).cuda(device).float()
    #B = Variable(torch.from_numpy( load_5D(opt.moving))).cuda(device).float()
    start2 = timeit.default_timer()
    print('Time for loading data: ', start2 - start)
    pred = model(A, B)
    F_AB = pred.permute(0, 2, 3, 4, 1).data.cpu().numpy()[0, :, :, :, :]
    #F_AB = pred.permute(0,2,3,4,1).data.cuda(device).numpy()[0, :, :, :, :]

    F_AB = F_AB.astype(np.float32) * range_flow
    warped_A = transform(A,
                         pred.permute(0, 2, 3, 4, 1) * range_flow,
                         grid).data.cpu().numpy()[0, 0, :, :, :]
    #warped_A = transform(A,pred.permute(0,2,3,4,1)*range_flow,grid).data.cuda(device).numpy()[0, 0, :, :, :]
    start3 = timeit.default_timer()
    print('Time for registration: ', start3 - start2)
    warped_F_BA = transform(-pred,
                            pred.permute(0, 2, 3, 4, 1) * range_flow,
                            grid).permute(0, 2, 3, 4,
                                          1).data.cpu().numpy()[0, :, :, :, :]
    #warped_F_BA = transform(-pred,pred.permute(0,2,3,4,1)*range_flow,grid).permute(0,2,3,4,1).data.cuda(device).numpy()[0, :, :, :, :]
    warped_F_BA = warped_F_BA.astype(np.float32) * range_flow
    start4 = timeit.default_timer()
    print('Time for generating inverse flow: ', start4 - start3)
    save_flow(F_AB, savepath + '/flow_A_B.nii.gz')
    save_flow(warped_F_BA, savepath + '/inverse_flow_B_A.nii.gz')
    save_img(warped_A, savepath + '/warped_A.nii.gz')
    start5 = timeit.default_timer()
    print('Time for saving results: ', start5 - start4)
    del pred
    pred = model(B, A)
    F_BA = pred.permute(0, 2, 3, 4, 1).data.cpu().numpy()[0, :, :, :, :]
    #F_BA = pred.permute(0,2,3,4,1).data.cuda(device).numpy()[0, :, :, :, :]
    F_BA = F_BA.astype(np.float32) * range_flow
    warped_B = transform(B,
                         pred.permute(0, 2, 3, 4, 1) * range_flow,
                         grid).data.cpu().numpy()[0, 0, :, :, :]
    #warped_B = transform(B,pred.permute(0,2,3,4,1)*range_flow,grid).data.cuda(device).numpy()[0, 0, :, :, :]
    warped_F_AB = transform(-pred,
                            pred.permute(0, 2, 3, 4, 1) * range_flow,
                            grid).permute(0, 2, 3, 4,
                                          1).data.cpu().numpy()[0, :, :, :, :]
    #warped_F_AB = transform(-pred,pred.permute(0,2,3,4,1)*range_flow,grid).permute(0,2,3,4,1).data.cuda(device).numpy()[0, :, :, :, :]
    warped_F_AB = warped_F_AB.astype(np.float32) * range_flow
    save_flow(F_BA, savepath + '/flow_B_A.nii.gz')
    save_flow(warped_F_AB, savepath + '/inverse_flow_A_B.nii.gz')
    save_img(warped_B, savepath + '/warped_B.nii.gz')