model/inference_demo.py

import argparse
import os
import torch
import torch.nn.functional as F
import json
from segment_anything_volumetric import sam_model_registry
from network.model import SegVol
from data_process.demo_data_process import process_ct_gt
import monai.transforms as transforms
from utils.monai_inferers_utils import sliding_window_inference, generate_box, select_points, build_binary_cube, build_binary_points, logits2roi_coor
from utils.visualize import draw_result

def set_parse():
    # %% set up parser
    parser = argparse.ArgumentParser()
    parser.add_argument("--test_mode", default=True, type=bool)
    parser.add_argument("--resume", type = str, default = '')
    parser.add_argument("-infer_overlap", default=0.5, type=float, help="sliding window inference overlap")
    parser.add_argument("-spatial_size", default=(32, 256, 256), type=tuple)
    parser.add_argument("-patch_size", default=(4, 16, 16), type=tuple)
    parser.add_argument('-work_dir', type=str, default='./work_dir')
    ### demo
    parser.add_argument('--demo_config', type=str, required=True)
    parser.add_argument("--clip_ckpt", type = str, default = './config/clip')
    args = parser.parse_args()
    return args

def dice_score(preds, labels):  # on GPU
    assert preds.shape[0] == labels.shape[0], "predict & target batch size don't match\n" + str(preds.shape) + str(labels.shape)
    predict = preds.view(1, -1)
    target = labels.view(1, -1)
    if target.shape[1] < 1e8:
        predict = predict.cuda()
        target = target.cuda()
    predict = torch.sigmoid(predict)
    predict = torch.where(predict > 0.5, 1., 0.)
    
    tp = torch.sum(torch.mul(predict, target))
    den = torch.sum(predict) + torch.sum(target) + 1
    dice = 2 * tp / den

    if target.shape[1] < 1e8:
        predict = predict.cpu()
        target = target.cpu()
    return dice

def zoom_in_zoom_out(args, segvol_model, image, image_resize, gt3D, gt3D_resize, categories=None):
    logits_labels_record = {}
    image_single_resize = image_resize
    image_single = image[0,0]
    ori_shape = image_single.shape
    for item_idx in range(len(categories)):
        # get label to generate prompts
        label_single = gt3D[0][item_idx]
        label_single_resize = gt3D_resize[0][item_idx]
        # skip meaningless categories
        if torch.sum(label_single) == 0:
            print('No object, skip')
            continue
        # generate prompts
        text_single = categories[item_idx] if args.use_text_prompt else None
        if categories is not None: print(f'inference |{categories[item_idx]}| target...')
        points_single = None
        box_single = None
        if args.use_point_prompt:
            point, point_label = select_points(label_single_resize, num_positive_extra=3, num_negative_extra=3)
            points_single = (point.unsqueeze(0).float().cuda(), point_label.unsqueeze(0).float().cuda()) 
            binary_points_resize = build_binary_points(point, point_label, label_single_resize.shape)
        if args.use_box_prompt:
            box_single = generate_box(label_single_resize).unsqueeze(0).float().cuda()
            binary_cube_resize = build_binary_cube(box_single, binary_cube_shape=label_single_resize.shape)
        
        ####################
        # zoom-out inference:
        print('--- zoom out inference ---')
        print(f'use text-prompt [{text_single!=None}], use box-prompt [{box_single!=None}], use point-prompt [{points_single!=None}]')
        with torch.no_grad():
            logits_global_single = segvol_model(image_single_resize.cuda(),
                                                text=text_single, 
                                                boxes=box_single, 
                                                points=points_single)
        
        # resize back global logits
        logits_global_single = F.interpolate(
                logits_global_single.cpu(),
                size=ori_shape, mode='nearest')[0][0]
        
        # build prompt reflection for zoom-in
        if args.use_point_prompt:
            binary_points = F.interpolate(
                binary_points_resize.unsqueeze(0).unsqueeze(0).float(),
                size=ori_shape, mode='nearest')[0][0]
        if args.use_box_prompt:
            binary_cube = F.interpolate(
                binary_cube_resize.unsqueeze(0).unsqueeze(0).float(),
                size=ori_shape, mode='nearest')[0][0]
        zoom_out_dice = dice_score(logits_global_single.squeeze(), label_single.squeeze())
        logits_labels_record[categories[item_idx]] = (
            zoom_out_dice,
            image_single, 
            points_single,
            box_single,
            logits_global_single, 
            label_single)
        print(f'zoom out inference done with zoom_out_dice: {zoom_out_dice:.4f}')
        if not args.use_zoom_in:
            continue

        ####################
        # zoom-in inference:
        min_d, min_h, min_w, max_d, max_h, max_w = logits2roi_coor(args.spatial_size, logits_global_single)
        if min_d is None:
            print('Fail to detect foreground!')
            continue

        # Crop roi
        image_single_cropped = image_single[min_d:max_d+1, min_h:max_h+1, min_w:max_w+1].unsqueeze(0).unsqueeze(0)
        global_preds = (torch.sigmoid(logits_global_single[min_d:max_d+1, min_h:max_h+1, min_w:max_w+1])>0.5).long()
        
        assert not (args.use_box_prompt and args.use_point_prompt)
        # label_single_cropped = label_single[min_d:max_d+1, min_h:max_h+1, min_w:max_w+1].unsqueeze(0).unsqueeze(0)
        prompt_reflection = None
        if args.use_box_prompt:
            binary_cube_cropped = binary_cube[min_d:max_d+1, min_h:max_h+1, min_w:max_w+1]
            prompt_reflection = (
                binary_cube_cropped.unsqueeze(0).unsqueeze(0),
                global_preds.unsqueeze(0).unsqueeze(0)
            )
        if args.use_point_prompt:
            binary_points_cropped = binary_points[min_d:max_d+1, min_h:max_h+1, min_w:max_w+1]
            prompt_reflection = (
                binary_points_cropped.unsqueeze(0).unsqueeze(0),
                global_preds.unsqueeze(0).unsqueeze(0)
            )
    
        ## inference
        with torch.no_grad():
            logits_single_cropped = sliding_window_inference(
                    image_single_cropped.cuda(), prompt_reflection,
                    args.spatial_size, 1, segvol_model, args.infer_overlap,
                    text=text_single,
                    use_box=args.use_box_prompt,
                    use_point=args.use_point_prompt,
                )
            logits_single_cropped = logits_single_cropped.cpu().squeeze()
        logits_global_single[min_d:max_d+1, min_h:max_h+1, min_w:max_w+1] = logits_single_cropped
        zoom_in_dice = dice_score(logits_global_single.squeeze(), label_single.squeeze())
        logits_labels_record[categories[item_idx]] = (
            zoom_in_dice,
            image_single, 
            points_single,
            box_single,
            logits_global_single, 
            label_single)
        print(f'===> zoom out dice {zoom_out_dice:.4f} -> zoom-out-zoom-in dice {zoom_in_dice:.4f} <===')
    return logits_labels_record

def inference_single_ct(args, segvol_model, data_item, categories):
    segvol_model.eval()
    image, gt3D = data_item["image"].float(), data_item["label"]
    image_zoom_out, gt3D__zoom_out = data_item["zoom_out_image"].float(), data_item['zoom_out_label']

    logits_labels_record = zoom_in_zoom_out(
        args, segvol_model, 
        image.unsqueeze(0), image_zoom_out.unsqueeze(0), 
        gt3D.unsqueeze(0), gt3D__zoom_out.unsqueeze(0),     # add batch dim
        categories=categories)
    
    # visualize
    if args.visualize:
        for target, values in logits_labels_record.items():
            dice_score, image, point_prompt, box_prompt, logits, labels = values
            print(f'{target} result with Dice score {dice_score:.4f} visualizing')
            draw_result(target + f"-Dice {dice_score:.4f}", image, box_prompt, point_prompt, logits, labels, args.spatial_size, args.work_dir)

def main(args):
    gpu = 0
    torch.cuda.set_device(gpu)
    # build model
    sam_model = sam_model_registry['vit'](args=args)
    segvol_model = SegVol(
                        image_encoder=sam_model.image_encoder, 
                        mask_decoder=sam_model.mask_decoder,
                        prompt_encoder=sam_model.prompt_encoder,
                        clip_ckpt=args.clip_ckpt,
                        roi_size=args.spatial_size,
                        patch_size=args.patch_size,
                        test_mode=args.test_mode,
                        ).cuda()
    segvol_model = torch.nn.DataParallel(segvol_model, device_ids=[gpu])

    # load param
    if os.path.isfile(args.resume):
        ## Map model to be loaded to specified single GPU
        loc = 'cuda:{}'.format(gpu)
        checkpoint = torch.load(args.resume, map_location=loc)
        segvol_model.load_state_dict(checkpoint['model'], strict=True)
        print("loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))

    # load demo config
    with open(args.demo_config, 'r') as file:
        config_dict = json.load(file)
    ct_path, gt_path, categories = config_dict['demo_case']['ct_path'], config_dict['demo_case']['gt_path'], config_dict['categories']

    # preprocess for data
    data_item = process_ct_gt(ct_path, gt_path, categories, args.spatial_size)   # keys: image, label

    # seg config for prompt & zoom-in-zoom-out
    args.use_zoom_in = True
    args.use_text_prompt = True
    args.use_box_prompt = True
    args.use_point_prompt = False
    args.visualize = False

    inference_single_ct(args, segvol_model, data_item, categories)

if __name__ == "__main__":
    args = set_parse()
    main(args)