def __getitem__(self, index):
name = self.imgs_list[index]
random.seed(0)
torch.manual_seed(0)
if self.prepare:
ori_img = Image.open(name)
img = Image.open(self.o_imgs_list[index])
img = np.asarray(img)
img = Image.fromarray(img)
img = self.transforms(img)
else:
ori_img = Image.open(name)
img = self.transforms(ori_img)
ori_label = Image.open(self.labels_list[index])
ori_label = np.asarray(ori_label)
label = 2 - ori_label // 127
if self.args.n_classes == 2: # Convert to binary map
label = np.where(label == 2, np.ones_like(label), np.zeros_like(label))
# Keep the same transformation
# TODO: There exists a bug, when you use ToTensor in transform.
# TODO: The value will be changed in label
random.seed(0)
torch.manual_seed(0)
label = np.asarray(self.gt_transforms(label.copy()))
label = torch.from_numpy(label.copy())
if not self.need_name:
if self.fine_tune:
return img, self.domain_label
return img, label
else:
ori_img = np.asarray(ori_img)
ori_label = np.expand_dims(ori_label, axis=-1)
return img, label, name, ori_img, ori_label
def label_onehotEncoding(label, num_class, backend='keras'):
"""
backend='pytorch'
#input shape (value 0,1,2,...) : (image_depth,image_height,image_width)
#output shape (values 0,1) : (num_class+1,image_depth,image_height,image_width)
#background is 0, so channel should be num_class+1
backend='keras'
#input shape (value 0,1,2,...) : (image_depth,image_height,image_width)
#output shape (values 0,1) : (num_class+1,image_depth,image_height,image_width)
#background is 0, so channel should be num_class+1
"""
dimension = len(label.shape)
if dimension != 3:
print('error')
else:
if backend == 'pytorch':
label_onehot = np.zeros((num_class, label.shape[0], label.shape[1],
label.shape[2])) #torch
else:
label_onehot = np.zeros(
(label.shape[0], label.shape[1], label.shape[2], num_class))
for idx in range(num_class):
label_temp = label.copy() # torch --> clone()
label_temp[label_temp != idx] = 0.
label_temp[label_temp != 0.] = 1.
if backend == 'pytorch':
label_onehot[idx] = label_temp
else:
label_onehot[..., idx] = label_temp
return label_onehot
def dilate_erode_labels(label):
"""Substract an eroded label to a dilated one in order to prevent
boundaries contact.
Parameters
----------
label : np.ndarray, np.uint or np.int
Labelled image with shape (y, x).
Returns
-------
label_final : np.ndarray, np.int64
Labelled image with shape (y, x).
"""
# check parameters
stack.check_array(label, ndim=2, dtype=[np.uint8, np.uint16, np.int64])
# handle 64 bit integer
if label.dtype == np.int64:
label = label.astype(np.uint16)
# erode-dilate mask
label_dilated = stack.dilation_filter(label, "disk", 2)
label_eroded = stack.erosion_filter(label, "disk", 2)
borders = label_dilated - label_eroded
label_final = label.copy()
label_final[borders > 0] = 0
label_final = label_final.astype(np.int64)
return label_final
def __getitem__(self, index):
if self.args.dataset == 'meta_ORIGA':
name = self.imgs_list[index + 101]
ori_label = loadmat(self.labels_list[index + 101])['mask']
else:
name = self.imgs_list[index]
ori_label = loadmat(self.labels_list[index])['mask']
ori_img = Image.open(name)
# Keep the same transformation
seed = np.random.randint(999999999)
random.seed(seed)
torch.manual_seed(seed)
# TODO: There exists a bug, when you use ToTensor in transform.
# TODO: The value will be changed in label
if self.args.dataset == 'meta_ORIGA':
ori_label = np.argmax(ori_label, axis=-1).astype(np.uint8).copy()
label = np.asarray(self.gt_transforms(ori_label.copy()))
#label = np.expand_dims(label, axis=-1)
label = torch.from_numpy(label.copy())
if self.args.n_classes == 2: # Convert to binary map
label = torch.where(label > 0,
torch.ones_like(label), torch.zeros_like(label))
if self.args.n_colors == 4:
meta_data = plt.imread(self.meta_list[index])
meta_data = np.expand_dims(meta_data, axis=-1)
random.seed(seed)
torch.manual_seed(seed)
img = self.transforms(ori_img)
if not self.need_name:
if self.args.n_colors == 4:
meta_data = np.transpose(meta_data, (2, 0, 1))
return img, meta_data, label
return img, label
else:
ori_img = cv2.resize(np.asarray(ori_img), (3072, 2048))
ori_label = cv2.resize(ori_label, (3072, 2048))
ori_label = np.expand_dims(ori_label, axis=-1)
return img, label, name, ori_img, ori_label
def main():
opt = TestOptions()
args = opt.initialize()
if not os.path.exists(args.save):
os.makedirs(args.save)
model = CreateSSLModel(args)
cudnn.enabled = True
cudnn.benchmark = True
model.train()
model.cuda()
targetloader = CreateTrgDataSSLLoader(args)
#predicted_label = np.zeros((len(targetloader), 1634, 1634))
predicted_label = np.zeros((len(targetloader), 2056, 2124))
#predicted_prob = np.zeros((len(targetloader), 1634, 1634))
predicted_prob = np.zeros((len(targetloader), 2056, 2124))
image_name = []
for index, batch in enumerate(targetloader):
if index % 10 == 0:
print('%d processd' % index)
image, _, name = batch
image = Variable(image).cuda()
_, output, _, _ = model(image, ssl=True)
output = nn.functional.softmax(output / args.alpha, dim=1)
#output = nn.functional.upsample(output, (1634, 1634), mode='bilinear', align_corners=True).cpu().data[0].numpy()
output = nn.functional.upsample(
output, (2056, 2124), mode='bilinear',
align_corners=True).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
label, prob = np.argmax(output, axis=2), np.max(output, axis=2)
cup_mask = get_bool(label, 0)
disc_mask = get_bool(label, 0) + get_bool(label, 1)
disc_mask = disc_mask.astype(np.uint8)
cup_mask = cup_mask.astype(np.uint8)
for i in range(5):
disc_mask = scipy.signal.medfilt2d(disc_mask, 19)
cup_mask = scipy.signal.medfilt2d(cup_mask, 19)
disc_mask = morphology.binary_erosion(disc_mask,
morphology.diamond(7)).astype(
np.uint8) # return 0,1
cup_mask = morphology.binary_erosion(cup_mask,
morphology.diamond(7)).astype(
np.uint8) # return 0,1
disc_mask = get_largest_fillhole(disc_mask)
cup_mask = get_largest_fillhole(cup_mask)
disc_mask = morphology.binary_dilation(disc_mask,
morphology.diamond(7)).astype(
np.uint8) # return 0,1
cup_mask = morphology.binary_dilation(cup_mask,
morphology.diamond(7)).astype(
np.uint8) # return 0,1
disc_mask = get_largest_fillhole(disc_mask).astype(
np.uint8) # return 0,1
cup_mask = get_largest_fillhole(cup_mask).astype(np.uint8)
label = disc_mask + cup_mask
predicted_label[index] = label.copy()
predicted_prob[index] = prob.copy()
image_name.append(name[0])
thres = []
for i in range(3):
x = predicted_prob[predicted_label == i]
if len(x) == 0:
thres.append(0)
continue
x = np.sort(x)
#print(len(x))
print(i, x)
print(np.sum(x < 0.5), '/', len(x), np.sum(x < 0.5) / len(x))
#thres.append(x[np.int(np.round(len(x) *int(args.p[i])))])
thres.append(0.5)
thres = np.array(thres)
print(thres)
# thres[thres > 0.6] = 0.6
color_labels = {0: 255, 1: 128, 2: 0}
#import pdb;pdb.set_trace()
for index in range(len(targetloader)):
name = image_name[index]
label = predicted_label[index] # label_min:0 label_max=2
#print(label)
prob = predicted_prob[index] # prob_min:0.5 prob_max=1.0
#print(prob)
for i, color in color_labels.items():
label[label == i] = color
#print(thres[i])
label[(prob > thres[i]) * (label == i)] = color
output = np.asarray(label, dtype=np.uint8)
output = Image.fromarray(output.astype(np.uint8))
#name = name.split('.')[0] + '.png'
name = name.split('_')[0] + '.bmp'
output.save('%s/%s' % (args.save, name))
disc_dice, cup_dice = calculate_dice(args.gt_dir, args.save,
args.devkit_dir)
print('===> disc_dice:' + str(round(disc_dice, 3)) + '\t' + 'cup_dice:' +
str(round(cup_dice, 3)))