neuronflow / blob_loss Goto Github PK

            # compute mean
            vprint("label_loss:", label_loss)
            # mean_label_loss = 0
            vprint("blobs in crop:", len(label_loss))
            if not len(label_loss) == 0:
                mean_label_loss = sum(label_loss) / len(label_loss)
                # mean_label_loss = sum(label_loss) / \
                #     torch.count_nonzero(label_loss)
                vprint("mean_label_loss", mean_label_loss)
                element_blob_loss.append(mean_label_loss)

	batch_length = multi_label.shape[0]
	element_blob_loss = []
	# loop over elements
	for element in range(batch_length):
		if element < batch_length:
			end_index = element + 1
		elif element == batch_length:
			end_index = None

		element_label = multi_label[element:end_index, ...]
		element_output = network_outputs[element:end_index, ...]
		from skimage import measure
		element_label = element_label.squeeze(0).cpu().numpy()
		element_label = element_label.squeeze(0)
		element_label = measure.label(element_label)
		element_label = torch.from_numpy(element_label)
		element_label = element_label.unsqueeze(0)
		element_label = element_label.unsqueeze(0).cuda()

		element_output = element_output.squeeze(0).detach().cpu().numpy()
		element_output = element_output.squeeze(0)
		element_output = measure.label(element_output)
		element_output = torch.from_numpy(element_output)
		element_output = element_output.unsqueeze(0)
		element_output = element_output.unsqueeze(0).cuda()
		# print('element_output', len(torch.unique(element_output)))
		# print('label',element_label.size())
		unique_labels = torch.unique(element_label)
		blob_count = len(unique_labels) - 1
		label_loss = []
		for ula in unique_labels:
			if ula != 0:
				label_mask = element_label > 0
				label_mask = ~label_mask
				label_mask[element_label == ula] = 1
				the_label = element_label == ula
				the_label_int = the_label.int()
				masked_output = element_output * label_mask
				# print('masked_output',masked_output.size(),'the_label_int',(the_label.float()).size())
				try:
					blob_loss = self.dice(masked_output, the_label_int)
					# print('bblos',blob_loss)
				except:
					# if int does not work we try float
					blob_loss = self.dice(masked_output, the_label.float())
					# print('bblos111', blob_loss)
				label_loss.append(blob_loss)
				# print('label_loss',label_loss)

			if not len(label_loss) == 0:
				mean_label_loss = sum(label_loss) / len(label_loss)
				element_blob_loss.append(mean_label_loss)
				# print('element_loss',element_blob_loss)

	if not len(element_blob_loss) == 0:
		mean_element_blob_loss = sum(element_blob_loss) / len(element_blob_loss)
	return mean_element_blob_loss