def show_side_results(dataset_name, dataset_version=0): network_strings = [ 'unet_2_channels__64_dims__wz_64', 'unet_2_channels__64_dims__wz_64_offset', 'METHOD_double' ] string_of_names = ['Gt', 'U Net 12 Embs', 'Multi Offset', 'Proposed'] data = {} counter = 1 for network in network_strings: directory = "results/" + dataset_name + "/" + "v_" + str( dataset_version) + "/" + network + "/mini_inst" print("reading " + directory) results = tensors_io.load_volume(directory) data[counter] = results counter += 1 data[0] = tensors_io.load_volume("results/" + dataset_name + "/" + "v_" + str(dataset_version) + "/gt") ''' tensors_io.save_subplots_compare(data[0], data[1], data[2], data[3], "RESULTS_COMPARISON", string_of_names=string_of_names) og = tensors_io.load_volume("results/" + dataset_name + "/" + "v_" + str(dataset_version) + "/original") og = torch.transpose(og, 3, 2) tensors_io.save_subvolume(torch.transpose(data[0], 3, 2), 'compare/gt') tensors_io.save_subvolume(og, 'compare/side_og') tensors_io.save_subvolume(torch.transpose(data[1], 3, 2), 'compare/side_unet') tensors_io.save_subvolume(torch.transpose(data[2], 3, 2), 'compare/side_unet_of') tensors_io.save_subvolume(torch.transpose(data[3], 3, 2), 'compare/side_proposed') ''' multi = tensors_io.load_volume( "results/" + dataset_name + "/" + "v_" + str(dataset_version) + "/unet_double_multi_2_channels__64_dims__wz_64/mini_inst") tensors_io.save_subvolume(torch.transpose(multi, 1, 3), 'compare/multi')
def recrop(params_t, resample_masks=0, factor=1, directory='.'): if (directory != '.'): path_of_interest = directory + '/subsampled' resample_masks = 2 masks = tensors_io.load_volume(directory, scale=2).unsqueeze(0) tensors_io.save_subvolume(masks, path_of_interest) exit() else: params_t.cleaning = False params_t.cleaning_sangids = False params_t.scale_p = factor if (resample_masks == 0): path_of_interest = "results/" + params_t.dataset_name + "/" + "v_" + str( params_t.dataset_version) + "/or_subsampled" elif (resample_masks == 1): path_of_interest = "results/" + params_t.dataset_name + "/" + "v_" + str( params_t.dataset_version) + "/gt_subsampled" elif (resample_masks == 2): if (params_t.debug_cluster_unet_double): params_t.network_string = "METHOD" path_of_interest = "results/" + params_t.dataset_name + "/" + "v_" + str( params_t.dataset_version ) + "/" + params_t.network_string + "/mini_inst_only_subsampled" params_t.testing_mask = "results/" + params_t.dataset_name + "/" + "v_" + str( params_t.dataset_version ) + "/" + params_t.network_string + "/mini_inst_only" else: path_of_interest = '' data_volume, masks, V_or = tensors_io.load_data_and_masks(params_t) path_temp = "" for local_path in path_of_interest.split("/"): path_temp = path_temp + local_path if not os.path.isdir(path_temp): os.mkdir(path_temp) path_temp = path_temp + "/" if (resample_masks == 0): tensors_io.save_subvolume(data_volume, path_of_interest) else: tensors_io.save_subvolume_instances(data_volume * 0, masks, path_of_interest)
def train_instance_segmentation_net(t_params, data_path_list, mask_path_list): print("Starting Training for instance segmentation for " + t_params.network_name) print("Window Size: {}. N Classes: {} N Embeddings: {} N Dim{}".format(t_params.cube_size, t_params.n_classes, t_params.n_embeddings, t_params.ndims)) net_i = t_params.net_i net_s = t_params.net criterion_i = t_params.criterion_i device = t_params.device data_volume_list = [] # Load Data Volume if(t_params.uint_16 is False): for data_path in data_path_list: data_volume_list.append((tensors_io.load_volume(data_path, scale=t_params.scale_p)).unsqueeze(0)) else: for data_path in data_path_list: data_volume_list.append((tensors_io.load_fibers_uint16(data_path, scale=t_params.scale_p)).unsqueeze(0)) if(t_params.cleaning_sangids is True): for i in range(len(data_volume_list)): data_volume_list[i][0, 0, ...] = tensors_io.clean_noise(data_volume_list[i][0, 0, ...], data_volume_list[i]) if(t_params.cleaning is True): for i in range(len(data_volume_list)): data_volume_list[i] = tensors_io.normalize_dataset(data_volume_list[i]) # Load Masks masks_list = [] for mask_path in mask_path_list: masks_list.append(tensors_io.load_volume_uint16(mask_path, scale=t_params.scale_p).long().unsqueeze(0)) # Optimizer and loss function optimizer = optim.Adam(net_i.parameters(), lr=t_params.lr) num_datasets = len(data_path_list) # Send the model to GPU net_i = net_i.to(device) net_s = net_s.to(device) for epoch in range(t_params.epochs_instance): print('Starting epoch {}/{}.'.format(epoch + 1, t_params.epochs_instance)) net_i.train() epoch_loss = 0 emb_total_loss = 0 seg_total_loss = 0 for i in range(t_params.batch_size + 30): data_volume = data_volume_list[i % num_datasets] masks = masks_list[i % num_datasets] (mini_V, mini_M) = tensors_io.random_crop_3D_image_batched(data_volume, masks, t_params.cube_size) mini_V = mini_V.to(device) true_masks = mini_M.to(device) # Evaluate Net embedding_output = net_i(mini_V) loss = criterion_i(embedding_output, true_masks, t_params) if(loss is None): continue optimizer.zero_grad() loss.backward() optimizer.step() epoch_loss += loss.item() emb_total_loss += loss.item() if((epoch) % 10 == 0 and epoch > 10): torch.save(net_i.state_dict(), t_params.net_weights_dir[1]) # quick_seg_inst_test(t_params, mini_V, mini_M) print('Dict Saved') print("loss: " + str(epoch_loss / i) + ", e_loss: " + str(emb_total_loss / i) + ", s_loss: " + str(seg_total_loss / i)) # save dictionary torch.save(net_i.state_dict(), t_params.net_weights_dir[1]) print("FINISHED TRAINING")
def train_r_net_embedded_offset(net, net_s, data_path, mask_path, data_list2, mask_list2, n_classes=2, n_embedded=16, cube_size=32, epochs=10, batch_size=1, Patches_per_Epoch=30, scale_p=2, pre_trained=False, device=None): print("Starting Offset Training...") GPU_YES = torch.cuda.is_available() if(device is None): device = torch.device("cuda:0" if GPU_YES else "cpu") if(pre_trained): print("Loading pre-trained-weights") net.load_state_dict(torch.load('info_files/r_net_offset.pth')) # Load Data Volume data_volume = tensors_io.load_volume(data_path, scale=scale_p) data_volume[0, ...] = tensors_io.clean_noise(data_volume[0, ...], data_path) data_volume = data_volume.unsqueeze(0) # Load Data Volume2 temp_volume = tensors_io.load_volume(data_list2[0], scale=scale_p) temp_volume[0, ...] = tensors_io.clean_noise(temp_volume[0, ...], data_list2[0]) masks = tensors_io.load_volume_uint16(mask_path, scale=scale_p).long().unsqueeze(0) (ch2, rows2, cols2, depth2) = temp_volume.shape num_datasets = len(data_list2) data_volume2 = torch.zeros(num_datasets, ch2, rows2, cols2, depth2) masks2 = torch.zeros(num_datasets, ch2, rows2, cols2, depth2, dtype=torch.long) data_volume2[0, ...] = temp_volume masks2[0, ...] = tensors_io.load_volume_uint16(mask_list2[0], scale=scale_p).long() for counter in range(1, num_datasets): temp_volume = tensors_io.load_volume(data_list2[counter], scale=scale_p) temp_volume[0, ...] = tensors_io.clean_noise(temp_volume[0, ...], data_list2[counter]) data_volume2[counter, ...] = temp_volume masks2[counter, ...] = tensors_io.load_volume_uint16(mask_list2[counter], scale=scale_p).long() [_, channels, rows, cols, slices] = data_volume.size() # Optimizer and loss function optimizer = optim.SGD(net.parameters(), lr=0.00001, momentum=0.9, weight_decay=0.0005) # Send the model to CPU or GPU net = net.to(device) for epoch in range(epochs): print('Starting epoch {}/{}.'.format(epoch + 1, epochs)) net.train() epoch_loss = 0 emb_total_loss = 0 seg_total_loss = 0 for i in range(Patches_per_Epoch): if(i % (num_datasets + 1) == 0): (mini_V, mini_M) = tensors_io.random_crop_3D_image_batched(data_volume, masks, cube_size) else: number = i % num_datasets vol = data_volume2[number, ...].unsqueeze(0) msk = masks2[number, ...].unsqueeze(0) (mini_V, mini_M) = tensors_io.random_crop_3D_image_batched(vol, msk, cube_size) mini_V = mini_V.to(device) true_masks = mini_M.to(device) # Evaluate Net embedding_output = net(mini_V) loss = embedded_geometric_loss_coords(embedding_output, true_masks) optimizer.zero_grad() loss.backward() optimizer.step() epoch_loss += loss.item() emb_total_loss += loss.item() if((epoch) % 10 == 0 and epoch > -1): torch.save(net.state_dict(), "info_files/r_net_offset.pth") # save_data(embedding_output, true_masks, epoch) print('Dict Saved') print("loss: " + str(epoch_loss / i) + ", e_loss: " + str(emb_total_loss / i) + ", s_loss: " + str(seg_total_loss / i)) # save dictionary torch.save(net.state_dict(), "info_files/r_net_offset.pth") print("FINISHED TRAINING")
def train_multitask_loss_net(t_params, data_path_list, mask_path_list): print("Starting Training for multitask loss with " + t_params.network_name + " in device " + str(t_params.device)) net_mt = t_params.net criterion_i = t_params.criterion_i criterion_s = t_params.criterion_s device = t_params.device data_volume_list = [] # Load Data Volume if(t_params.uint_16 is False): for data_path in data_path_list: data_volume_list.append((tensors_io.load_volume(data_path, scale=t_params.scale_p)).unsqueeze(0)) else: for data_path in data_path_list: data_volume_list.append((tensors_io.load_fibers_uint16(data_path, scale=t_params.scale_p)).unsqueeze(0)) if(t_params.cleaning_sangids is True): for i in range(len(data_volume_list)): data_volume_list[i][0, 0, ...] = tensors_io.clean_noise(data_volume_list[i][0, 0, ...], data_volume_list[i]) if(t_params.cleaning is True): for i in range(len(data_volume_list)): data_volume_list[i] = tensors_io.normalize_dataset(data_volume_list[i]) # Load Masks masks_list = [] for mask_path in mask_path_list: masks_list.append(tensors_io.load_volume_uint16(mask_path, scale=t_params.scale_p).long().unsqueeze(0)) # Optimizer and loss function optimizer = optim.Adam(net_mt.parameters(), lr=t_params.lr) num_datasets = len(data_path_list) # Send the model to GPU net_mt = net_mt.to(device) for epoch in range(t_params.epochs_instance): print('Starting epoch {}/{}.'.format(epoch + 1, t_params.epochs_instance)) net_mt.train() epoch_loss = 0 emb_total_loss = 0 seg_total_loss = 0 for i in range(t_params.batch_size + 30): data_volume = data_volume_list[i % num_datasets] masks = masks_list[i % num_datasets] (mini_V, mini_M) = tensors_io.random_crop_3D_image_batched(data_volume, masks, t_params.cube_size) mini_V = mini_V.to(device) true_masks = mini_M.to(device) if(true_masks.max() == 0): sigma0 = torch.tensor(1) sigma1 = torch.tensor(1) continue # Evaluate Net outputs = net_mt(mini_V) if(t_params.network_name == "unet_double"): segmentation_output = outputs[0] embedding_output = outputs[1] s_loss = criterion_s((true_masks > 0).long(), segmentation_output) e_loss = criterion_i(embedding_output, true_masks, t_params) if(e_loss is None or s_loss is None): continue total_loss = t_params.alpha_seg * s_loss + t_params.alpha_emb * e_loss sigma0 = torch.tensor(t_params.alpha_seg) sigma1 = torch.tensor(t_params.alpha_emb) else: total_loss, s_loss, e_loss, sigma0, sigma1 = t_params.criterion(outputs, true_masks, t_params) optimizer.zero_grad() total_loss.backward() optimizer.step() epoch_loss += total_loss.item() emb_total_loss += e_loss.item() seg_total_loss += s_loss.item() if((epoch) % 10 == 0): t_params.training_numbers['total_loss'].append(epoch_loss) t_params.training_numbers['seg_loss'].append(seg_total_loss) t_params.training_numbers['offset_loss'].append(emb_total_loss) t_params.training_numbers['sigma0'].append(sigma0.cpu().item()) t_params.training_numbers['sigma1'].append(sigma1.cpu().item()) torch.save(net_mt.state_dict(), t_params.net_weights_dir[0]) np.save(t_params.net_train_dict_dir, t_params.training_numbers) # mini_V_or, final_pred, final_fibers, mini_M, seg_eval, inst_eval, inst_eval_objectwise = quick_seg_inst_test(t_params) # t_params.training_numbers['seg_f1_train'].append(seg_eval) # t_params.training_numbers['Ra_train'].append(inst_eval) print('Dict Saved') print("loss: " + str(epoch_loss / i) + ", e_loss: " + str(emb_total_loss / i) + ", s_loss: " + str(seg_total_loss / i) + ", sigma0:" + str(sigma0.cpu().item()) + ", sigma1:" + str(sigma1.cpu().item())) # save dictionary torch.save(net_mt.state_dict(), t_params.net_weights_dir[0]) print("FINISHED TRAINING")
def train_semantic_segmentation_net(t_params, data_path_list, mask_path_list): print("Starting Training for Semantic Segmentation: " + t_params.network_name) net = t_params.net criterion = t_params.criterion_s device = t_params.device data_volume_list = [] # Load Data Volume if(t_params.uint_16 is False): for data_path in data_path_list: data_volume_list.append((tensors_io.load_volume(data_path, scale=t_params.scale_p)).unsqueeze(0)) else: for data_path in data_path_list: data_volume_list.append((tensors_io.load_fibers_uint16(data_path, scale=t_params.scale_p)).unsqueeze(0)) if(t_params.cleaning_sangids is True): for i in range(len(data_volume_list)): data_volume_list[i][0, 0, ...] = tensors_io.clean_noise(data_volume_list[i][0, 0, ...], data_volume_list[i]) if(t_params.cleaning is True): for i in range(len(data_volume_list)): data_volume_list[i] = tensors_io.normalize_dataset(data_volume_list[i]) # Load Masks masks_list = [] for mask_path in mask_path_list: masks_list.append(tensors_io.load_volume_uint16(mask_path, scale=t_params.scale_p).long().unsqueeze(0)) # Optimizer and loss function optimizer = optim.Adam(net.parameters(), lr=t_params.lr) num_datasets = len(data_path_list) # Send the model to GPU print("Training for {} epochs".format(t_params.epochs)) net = net.to(device) for epoch in range(t_params.epochs): print('Starting epoch {}/{}.'.format(epoch + 1, t_params.epochs)) net.train() epoch_loss = 0 seg_total_loss = 0 for i in range(t_params.batch_size + 30): data_volume = data_volume_list[i % num_datasets] masks = masks_list[i % num_datasets] (mini_V, mini_M) = tensors_io.random_crop_3D_image_batched(data_volume, masks, t_params.cube_size, random_rotations=1) mini_V = mini_V.to(device) mini_M = mini_M.to(device) if(t_params.n_classes == 2): true_masks = (mini_M > 0).long() else: true_masks = mini_M.long() if(true_masks.max() == 0): continue segmentation_output = net(mini_V) s_loss = criterion(true_masks, segmentation_output) epoch_loss += s_loss.item() seg_total_loss += s_loss.item() optimizer.zero_grad() s_loss.backward() optimizer.step() if(epoch % 10 == 0): print('Dict Saved') torch.save(net.state_dict(), t_params.net_weights_dir[0]) _, final_pred = segmentation_output.max(1) evaluation.evaluate_segmentation(final_pred, mini_M) if(t_params.debug is True): tensors_io.save_subvolume_instances(mini_V, final_pred, "results/debug_seg_training") print("loss: " + str(epoch_loss / i)) if(epoch_loss / i < 0.01): break # save dictionary torch.save(net.state_dict(), t_params.net_weights_dir[0]) evaluation.evaluate_segmentation(final_pred, mini_M) print("FINISHED TRAINING")