def predict_central_focus(image, out_file): in_channels = 3 out_channels = 1 init_features = 32 image_size = 224 model_path = "./pytorch_unet_models/unet_tumors_only2.pt" preprocess = [ resample, stack_channels_valid, normalization, pad_and_resize ] device = torch.device("cpu" if not torch.cuda.is_available() else "cuda:0") transform = transforms.Compose([ transforms.Resize(image_size), transforms.ToTensor(), ]) unet = UNet(in_channels=in_channels, out_channels=out_channels, init_features=init_features) unet.to(device) unet.load_state_dict(torch.load(model_path)) unet.eval() epi_image = nib.load(image) epi_image_data = epi_image.get_fdata() dummy = epi_image_data.copy() h, w = epi_image_data.shape[0], epi_image_data.shape[1] pixdims = (epi_image.header["pixdim"], epi_image.header["pixdim"]) for preprocess_ in preprocess: if preprocess_ == pad_and_resize: epi_image_data, dummy = preprocess_(epi_image_data, dummy, image_size=image_size) elif preprocess_ == resample: epi_image_data, dummy = preprocess_(epi_image_data, dummy, pixdims=pixdims) else: epi_image_data, dummy = preprocess_(epi_image_data, dummy) epi_label_pred = [] for n_slice in trange(epi_image_data.shape[-1]): if n_slice > epi_image_data.shape[ -1] * 0.05 and n_slice < epi_image_data.shape[-1] * 0.95: input_image = epi_image_data[..., n_slice] x = transform(Image.fromarray(input_image.transpose(1, 0, 2))) x = torch.unsqueeze(x, 0) y_pred = unet(x.to(device)) y_pred_np = torch.squeeze(y_pred).detach().cpu().numpy() y_pred_np = y_pred_np.transpose(1, 0) y_pred_np = cv2.resize(y_pred_np, (h, w), cv2.INTER_CUBIC) #y_pred_np = np.round(y_pred_np).astype(np.uint8) else: y_pred_np = np.zeros((h, w), dtype=np.uint8) epi_label_pred.append(np.expand_dims(y_pred_np, axis=-1)) epi_label_pred = np.concatenate(epi_label_pred, axis=-1) epi_label_pred = np.round(epi_label_pred) img = nib.Nifti1Image(epi_label_pred, affine=None) img.to_filename(out_file)
"2-stage-256-2d-unet64": 0.5316030 * 0.2, "2-stage-256-2d-rc-unet64": 0.4642530 * 0.2, "2-stage-256-2d-n4rc-unet64": 0.4782162 * 0.2, "2-stage-224-2d-n4-unet32": 0.4700288 * 0.2, "128-3d-unet16": 0.5915776 * 0.5, "128-3d-unet16-pretrained": 0.6051299 * 0.5 } model_list = os.listdir(feature_dir) device = torch.device("cpu" if not torch.cuda.is_available() else "cuda:0") in_channels = len(model_list) unet = UNet(in_channels=in_channels, out_channels=out_channels, init_features=init_features) unet.load_state_dict(torch.load(weights)) unet.to(device, dtype=torch.float) unet.eval() filenames = glob.glob(os.path.join(feature_dir, model_list[0], "*.nii.gz")) filenames = [os.path.basename(filename) for filename in filenames] for filename in tqdm(filenames): feature_paths = {} for model in model_list: feature_paths[model] = os.path.join(feature_dir, model, filename) features = {key: nib.load(value).get_fdata() for key, value in feature_paths.items()} y_pred = []
def predict(image, out_file): in_channels = 3 out_channels = 1 init_features = 64 origin_image_size = 512 image_size = 256 labels_map = {"brain": 0, "nobrain": 1} unet_model_path = "./pytorch_unet_models/rc-unet64.pt" efficientnet_model_path = "./pytorch_efficientnet_models/nobrainer.pt" model_name = "efficientnet-b0" preprocess = [ resample, stack_channels_valid, normalization, pad_and_resize ] device = torch.device("cpu" if not torch.cuda.is_available() else "cuda:0") #device = torch.device("cpu") nobrainer_transform = transforms.Compose([ transforms.Resize(image_size), transforms.ToTensor(), ]) unet_transform = transforms.Compose([ transforms.FiveCrop(image_size), transforms.Lambda(lambda crops: torch.stack( [transforms.ToTensor()(crop) for crop in crops])) ]) #efn = EfficientNet.from_name(model_name, num_classes=len(labels_map)) efn = EfficientNet.from_name( model_name, override_params={'num_classes': len(labels_map)}) efn.load_state_dict(torch.load(efficientnet_model_path)) efn.to(device) unet = UNet(in_channels=in_channels, out_channels=out_channels, init_features=init_features) unet.to(device) unet.load_state_dict(torch.load(unet_model_path)) unet.eval() epi_image = nib.load(image) epi_image_data = epi_image.get_fdata() dummy = epi_image_data.copy() h, w = epi_image_data.shape[0], epi_image_data.shape[1] pixdims = (epi_image.header["pixdim"], epi_image.header["pixdim"]) for preprocess_ in preprocess: if preprocess_ == pad_and_resize: epi_image_data, dummy = preprocess_(epi_image_data, dummy, image_size=origin_image_size) elif preprocess_ == resample: epi_image_data, dummy = preprocess_(epi_image_data, dummy, pixdims=pixdims) else: epi_image_data, dummy = preprocess_(epi_image_data, dummy) epi_label_pred = [] no_brain = [] # stage one for n_slice in trange(epi_image_data.shape[-1]): input_image = epi_image_data[..., n_slice] x = nobrainer_transform(Image.fromarray(input_image.transpose(1, 0, 2))) x = torch.unsqueeze(x, 0) y_nobrain = efn(x.to(device)) y_nobrain_np = torch.squeeze(y_nobrain).detach().cpu().numpy() no_brain.append(np.argmax(softmax(y_nobrain_np))) # expand from center n_center_slice = int(epi_image_data.shape[-1] / 2) for n_slice in range(1, n_center_slice): prev, current = no_brain[n_slice - 1], no_brain[n_slice] if current == labels_map["brain"] and prev == labels_map["nobrain"]: no_brain[n_slice:n_center_slice] = [labels_map["brain"] ] * (n_center_slice - n_slice) break for n_slice in range(n_center_slice, epi_image_data.shape[-1]): prev, current = no_brain[n_slice - 1], no_brain[n_slice] if current == labels_map["nobrain"] and prev == labels_map["brain"]: no_brain[n_slice:epi_image_data.shape[-1]] = [ labels_map["nobrain"] ] * (epi_image_data.shape[-1] - n_slice) break # stage two for n_slice in trange(epi_image_data.shape[-1]): if no_brain[n_slice] == labels_map["brain"]: input_image = epi_image_data[..., n_slice] x = unet_transform(Image.fromarray(input_image.transpose(1, 0, 2))) y_pred = unet(x.to(device)) y_pred_np = torch.squeeze(y_pred).detach().cpu().numpy() y_pred_np = reconstruct_label(y_pred_np, h, w) y_pred_np = y_pred_np.transpose(1, 0) #y_pred_np = np.round(y_pred_np).astype(np.uint8) elif no_brain[n_slice] == labels_map["nobrain"]: y_pred_np = np.zeros((h, w), dtype=np.uint8) epi_label_pred.append(np.expand_dims(y_pred_np, axis=-1)) epi_label_pred = np.concatenate(epi_label_pred, axis=-1) epi_label_pred = np.round(epi_label_pred) print("nobrain count: {:d}".format(np.sum(no_brain))) img = nib.Nifti1Image(epi_label_pred, affine=None) img.to_filename(out_file)
def predict(image, out_file): in_channels = 3 out_channels = 1 init_features = 64 origin_image_size = 512 image_size = 256 labels_map = {"brain": 0, "nobrain": 1} unet_model_path = "../weights/unet64/unet_last.pt" preprocess = [ resample, stack_channels_valid, normalization, pad_and_resize ] device = torch.device("cpu" if not torch.cuda.is_available() else "cuda:0") #device = torch.device("cpu") unet_transform = transforms.Compose([ transforms.FiveCrop(image_size), transforms.Lambda(lambda crops: torch.stack( [transforms.ToTensor()(crop) for crop in crops])) ]) unet = UNet(in_channels=in_channels, out_channels=out_channels, init_features=init_features) unet.to(device) unet.load_state_dict(torch.load(unet_model_path)) unet.eval() epi_image = nib.load(image) epi_image_data = epi_image.get_fdata() dummy = epi_image_data.copy() h, w = epi_image_data.shape[0], epi_image_data.shape[1] pixdims = (epi_image.header["pixdim"], epi_image.header["pixdim"]) for preprocess_ in preprocess: if preprocess_ == pad_and_resize: epi_image_data, dummy = preprocess_(epi_image_data, dummy, image_size=origin_image_size) elif preprocess_ == resample: epi_image_data, dummy = preprocess_(epi_image_data, dummy, pixdims=pixdims) else: epi_image_data, dummy = preprocess_(epi_image_data, dummy) epi_label_pred = [] # stage two for n_slice in trange(epi_image_data.shape[-1]): input_image = epi_image_data[..., n_slice] x = unet_transform(Image.fromarray(input_image.transpose(1, 0, 2))) y_pred = unet(x.to(device)) y_pred_np = torch.squeeze(y_pred).detach().cpu().numpy() y_pred_np = reconstruct_label(y_pred_np, h, w) y_pred_np = y_pred_np.transpose(1, 0) y_pred_np = np.round(y_pred_np).astype(np.uint8) epi_label_pred.append(np.expand_dims(y_pred_np, axis=-1)) epi_label_pred = np.concatenate(epi_label_pred, axis=-1) img = nib.Nifti1Image(epi_label_pred, affine=None) img.to_filename(out_file)
os.makedirs(fig_dir) image_size = 224 device = torch.device("cpu" if not torch.cuda.is_available() else "cuda:0") transforms = transforms.Compose([ transforms.Resize(image_size), transforms.ToTensor(), ]) valid_dataset = SegmentationDataset(valid_folder_path, transform=transforms) unet = UNet(in_channels=3, out_channels=1) unet.to(device) unet.load_state_dict(torch.load("unet.pt")) unet.eval() count = 0 #for x,y in tqdm(valid_dataset): for x,y in valid_dataset: x = torch.unsqueeze(x, 0) y_pred = unet(x.to(device)) y_pred_np = torch.squeeze(y_pred).detach().cpu().numpy() y_true_np = torch.squeeze(y).detach().cpu().numpy() x_np = torch.squeeze(x).detach().cpu().numpy() #print(np.unique(y_true_np), (np.min(y_pred_np), np.max(y_pred_np))) x_np = x_np.transpose(1,2,0) #print(x_np.shape)
transform = transforms.Compose([ transforms.Resize(image_size), transforms.ToTensor(), ]) efn = EfficientNet.from_name( model_name, override_params={'num_classes': len(labels_map)}) efn.load_state_dict(torch.load(efficientnet_model_path)) efn.to(device) unet = UNet(in_channels=in_channels, out_channels=out_channels, init_features=init_features) unet.to(device) unet.load_state_dict(torch.load(unet_model_path)) unet.eval() epi_image = nib.load(image) epi_image_data = epi_image.get_fdata() dummy = epi_image_data.copy() h, w = epi_image_data.shape[0], epi_image_data.shape[1] pixdims = (epi_image.header["pixdim"], epi_image.header["pixdim"]) for preprocess_ in preprocess: if preprocess_ == pad_and_resize: epi_image_data, dummy = preprocess_(epi_image_data, dummy, image_size=image_size) elif preprocess_ == resample: epi_image_data, dummy = preprocess_(epi_image_data,