import argparse import os import torch import torch.nn.functional as F import json import monai.transforms as transforms from model.segment_anything_volumetric import sam_model_registry from model.network.model import SegVol from model.data_process.demo_data_process import process_ct_gt from model.utils.monai_inferers_utils import sliding_window_inference, generate_box, select_points, build_binary_cube, build_binary_points, logits2roi_coor from model.utils.visualize import draw_result import streamlit as st 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 = 'SegVol_v1.pth') parser.add_argument("-infer_overlap", default=0.0, 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("--clip_ckpt", type = str, default = 'model/config/clip') args = parser.parse_args() return args def zoom_in_zoom_out(args, segvol_model, image, image_resize, text_prompt, point_prompt, box_prompt): image_single_resize = image_resize image_single = image[0,0] ori_shape = image_single.shape resize_shape = image_single_resize.shape[2:] # generate prompts text_single = None if text_prompt is None else [text_prompt] points_single = None box_single = None if args.use_point_prompt: point, point_label = point_prompt points_single = (point.unsqueeze(0).float(), point_label.unsqueeze(0).float()) binary_points_resize = build_binary_points(point, point_label, resize_shape) if args.use_box_prompt: box_single = box_prompt.unsqueeze(0).float() binary_cube_resize = build_binary_cube(box_single, binary_cube_shape=resize_shape) #################### # zoom-out inference: print('--- zoom out inference ---') print(text_single) 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, 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] # draw_result('unknow', image_single_resize, None, point_prompt, logits_global_single, logits_global_single) if not args.use_zoom_in: return logits_global_single #################### # 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!') return logits_global_single # 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, 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_global_single=logits_global_single, ) logits_single_cropped = logits_single_cropped.cpu().squeeze() if logits_single_cropped.shape != logits_global_single.shape: logits_global_single[min_d:max_d+1, min_h:max_h+1, min_w:max_w+1] = logits_single_cropped return logits_global_single @st.cache_resource def build_model(): # build model st.write('building model') clip_ckpt = 'model/config/clip' resume = 'SegVol_v1.pth' sam_model = sam_model_registry['vit']() segvol_model = SegVol( image_encoder=sam_model.image_encoder, mask_decoder=sam_model.mask_decoder, prompt_encoder=sam_model.prompt_encoder, clip_ckpt=clip_ckpt, roi_size=(32,256,256), patch_size=(4,16,16), test_mode=True, ) segvol_model = torch.nn.DataParallel(segvol_model) segvol_model.eval() # load param if os.path.isfile(resume): ## Map model to be loaded to specified single GPU loc = 'cpu' checkpoint = torch.load(resume, map_location=loc) segvol_model.load_state_dict(checkpoint['model'], strict=True) print("loaded checkpoint '{}' (epoch {})".format(resume, checkpoint['epoch'])) print('model build done!') return segvol_model @st.cache_data def inference_case(_image, _image_zoom_out, _point_prompt, text_prompt, _box_prompt): # seg config args = set_parse() args.use_zoom_in = True args.use_text_prompt = text_prompt is not None args.use_box_prompt = _box_prompt is not None args.use_point_prompt = _point_prompt is not None segvol_model = build_model() # run inference logits = zoom_in_zoom_out( args, segvol_model, _image.unsqueeze(0), _image_zoom_out.unsqueeze(0), text_prompt, _point_prompt, _box_prompt) print(logits.shape) resize_transform = transforms.Compose([ transforms.AddChannel(), transforms.Resize((325,325,325), mode='trilinear') ] ) logits = resize_transform(logits)[0] print(logits.shape) return (torch.sigmoid(logits) > 0.5).int().numpy()