def SavePatchMap(MatPath, FigPath, MatName):
FigName = MatName.replace('.npz', '.png')
FigFile = os.path.join(FigPath, FigName)
MatFile = os.path.join(MatPath, MatName)
SegRes = load_npz(MatFile)
SegRes = SegRes.todense()
SegRes = vl2im(SegRes)
Fig = Image.fromarray(SegRes.astype(dtype=np.uint8))
del SegRes
Fig.convert('RGB')
Fig.save(FigFile, 'PNG')
def postprocess_rmsl(predicted_image_dir_path, new_image_path, ref_area,
ref_ecc):
if not os.path.exists(new_image_path):
os.makedirs(new_image_path)
for pth, dir, filenames in os.walk(predicted_image_dir_path):
for filename in filenames:
img_pth = os.path.join(pth, filename)
img = imread(img_pth)
img = im2vl(img)
img_close = closing(img, square(3))
labeled_img = label(img_close, connectivity=2)
new_labeled_img = remove_small_objects(labeled_img,
min_size=ref_area,
connectivity=2)
# remove non-circle region
props = regionprops(new_labeled_img)
for i, reg in enumerate(props):
if reg.eccentricity > ref_ecc:
new_labeled_img[new_labeled_img == reg.label] = 0
new_img = np.asarray(new_labeled_img != 0, dtype=np.uint8)
new_img = vl2im(new_img)
new_name = os.path.join(new_image_path, filename)
imsave(new_name, new_img)
def main():
"""Create the model and start the training."""
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
cudnn.enabled = True
torch.manual_seed(args.random_seed)
model = UNet(args.num_classes)
model = torch.nn.DataParallel(model)
optimizer = optim.SGD(model.parameters(),
lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay)
if args.fine_tune:
saved_state_dict = torch.load(args.restore_from)
saved_state_dict = saved_state_dict['state_dict']
new_params = model.state_dict().copy()
for i in saved_state_dict:
i_parts = i.split('.')
if 'final' not in i_parts:
new_params[i] = saved_state_dict[i]
model.load_state_dict(new_params)
elif args.restore_from:
if os.path.isfile(args.restore_from):
print("=> loading checkpoint '{}'".format(args.restore_from))
checkpoint = torch.load(args.restore_from)
try:
if args.set_start:
args.start_step = int(math.ceil(checkpoint['example'] / args.batch_size))
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (step {})"
.format(args.restore_from, args.start_step))
except:
model.load_state_dict(checkpoint)
print("=> loaded checkpoint '{}'".format(args.restore_from))
else:
print("=> no checkpoint found at '{}'".format(args.restore_from))
exit(0)
if not args.is_training:
# Frozen BN
# when training, the model will use the running means and the
# running vars of the pretrained model.
# But note that eval() doesn't turn off history tracking.
print("Freezing BN layers")
model.eval()
else:
print("Normal BN layers")
model.train()
model.cuda()
cudnn.benchmark = True
trainloader = data.DataLoader(LEAFTrain(args.data_list,
crop_size=input_size, scale=args.random_scale,
mirror=args.random_mirror, color_jitter=args.random_jitter,
rotate=args.random_rotate,
mean=IMG_MEAN),
batch_size=args.batch_size,
shuffle=True, num_workers=args.num_workers)
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
accuracy = AverageMeter()
writer = SummaryWriter(args.snapshot_dir)
# weight computation
class_distri = np.array(args.class_distri)
normalized_class_distri = class_distri/np.sum(class_distri)
class_weight = 1 / normalized_class_distri
cnt = 0
actual_step = args.start_step
while actual_step < args.final_step:
iter_end = timeit.default_timer()
for i_iter, batch in enumerate(trainloader):
actual_step = int(args.start_step + cnt)
data_time.update(timeit.default_timer() - iter_end)
images, labels, patch_name = batch
optimizer.zero_grad()
adjust_learning_rate(optimizer, actual_step)
pred = model(images)
image = images.data.cpu().numpy()[0]
labels = resize_target(labels, pred.size(2))
loss = loss_calc(pred, labels, class_weight)
losses.update(loss.item(), pred.size(0))
acc = _pixel_accuracy(pred.data.cpu().numpy(), labels.data.cpu().numpy())
accuracy.update(acc, pred.size(0))
loss.backward()
optimizer.step()
batch_time.update(timeit.default_timer() - iter_end)
iter_end = timeit.default_timer()
if actual_step % args.print_freq == 0:
print('iter: [{0}]{1}/{2}\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Pixel Accuracy {accuracy.val:.3f} ({accuracy.avg:.3f})'.format(
cnt, actual_step, args.final_step, batch_time=batch_time,
data_time=data_time, loss=losses, accuracy=accuracy))
writer.add_scalar("train_loss", losses.avg, actual_step)
writer.add_scalar("pixel_accuracy", accuracy.avg, actual_step)
writer.add_scalar("lr", optimizer.param_groups[0]['lr'], actual_step)
if actual_step % args.save_img_freq == 0:
msk_size = pred.size(2)
image = image.transpose(1, 2, 0)
image = cv2.resize(image, (msk_size, msk_size), interpolation=cv2.INTER_NEAREST)
image += IMG_MEAN
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
label = labels.data.cpu().numpy()[0]
label = vl2im(label)
single_pred = pred.data.cpu().numpy()[0].argmax(axis=0)
single_pred = vl2im(single_pred)
new_im = Image.new('RGB', (msk_size * 3, msk_size))
new_im.paste(Image.fromarray(image.astype('uint8'), 'RGB'), (0, 0))
new_im.paste(Image.fromarray(single_pred.astype('uint8'), 'RGB'), (msk_size, 0))
new_im.paste(Image.fromarray(label.astype('uint8'), 'RGB'), (msk_size * 2, 0))
new_im_name = 'B' + format(args.batch_size, "04d") + '_S' + format(actual_step, "06d") + '_' + patch_name[0]
new_im_file = os.path.join(args.img_dir, new_im_name)
new_im.save(new_im_file)
if actual_step % args.save_pred_every == 0 and cnt != 0:
print('taking snapshot ...')
torch.save({'example': actual_step * args.batch_size,
'state_dict': model.state_dict()},
osp.join(args.snapshot_dir,
'LEAF_UNET_B' + format(args.batch_size, "04d") + '_S' + format(actual_step, "06d") + '.pth'))
if actual_step >= args.final_step:
break
cnt += 1
print('save the final model ...')
torch.save({'example': actual_step * args.batch_size,
'state_dict': model.state_dict()},
osp.join(args.snapshot_dir,
'LEAF_UNET_B' + format(args.batch_size, "04d") + '_S' + format(actual_step, "06d") + '.pth'))
end = timeit.default_timer()
print(end - start, 'seconds')
def process(self, n):
os.environ["CUDA_VISIBLE_DEVICES"] = n
while imagelst:
# input
device = torch.device(
'cuda:0' if torch.cuda.is_available() else 'cpu')
tumorname = "green"
test_img_pth = imagelst[0]
# setting
patch_size = 512
INPUT_SIZE = (patch_size, patch_size)
overlap_size = 64
ref_area = 10000 # pre-processing
ref_extent = 0.6 # pre-processing
rm_rf_area = 5000 # post-processing
ref_ecc = 0.92 # post-processing
BATCH_SIZE = 32
model_id = '21700' # 'BACH_UNET_B2_S110000_Frozen_BN_test2048'
RESTORE_FROM = 'LEAF_UNET_B0064_S021700.pth'
root_pth = 'result' # may change to './' in the final version
savepath = root_pth + tumorname + '/512_test_stride_64'
logdirpth = root_pth + tumorname + '/log'
if not os.path.exists(logdirpth):
os.makedirs(logdirpth)
logpath = root_pth + tumorname + '/log/512_test_stride_64_XY3c.log'
images = os.path.join(savepath, 'images')
if not os.path.exists(images):
os.makedirs(images)
log = open(logpath, 'w')
log.write(test_img_pth + '\n')
filename = test_img_pth.split('/')[-1]
slidename = filename[:-4]
# divide testing image
process_tif(test_img_pth,
filename,
images,
log,
patch_size,
overlap_size,
ref_area=ref_area,
ref_extent=ref_extent)
# make testing dataset
txtdirpth = os.path.join(root_pth, tumorname, 'txt')
if not os.path.exists(txtdirpth):
os.makedirs(txtdirpth)
txtname = slidename + ".txt"
txtfile = os.path.join(txtdirpth, txtname)
txt = open(txtfile, 'w')
SlideDir = os.path.join(images, slidename)
PatchList = os.listdir(SlideDir)
for PatchName in PatchList:
PatchFile = os.path.join(SlideDir, PatchName)
txt.write(PatchFile + '\n')
txt.close()
#################################################
# testing
IMG_MEAN = np.array(
(62.17962105572224, 100.62603236734867, 131.60830906033516),
dtype=np.float32)
DATA_DIRECTORY = images
DATA_LIST_PATH = txtdirpth
NUM_CLASSES = 2
NPZ_PATH = root_pth + tumorname + '/PatchNet/npz/' + model_id
if not os.path.exists(NPZ_PATH):
os.makedirs(NPZ_PATH)
MAP_PATH = root_pth + tumorname + '/PatchNet/map/' + model_id
if not os.path.exists(MAP_PATH):
os.makedirs(MAP_PATH)
LOG_PATH = root_pth + tumorname + '/PatchNet/logfiles'
if not os.path.exists(LOG_PATH):
os.makedirs(LOG_PATH)
preprocess_start_time = time.time()
print("Restored from:", RESTORE_FROM)
LogName = "Test_HeatMap_log.txt"
LogFile = os.path.join(LOG_PATH, LogName)
log = open(LogFile, 'w')
log.writelines('batch size:' + str(BATCH_SIZE) + '\n')
log.writelines(DATA_LIST_PATH + '\n')
log.writelines('restore from ' + RESTORE_FROM + '\n')
model = UNet(NUM_CLASSES)
model = nn.DataParallel(model)
model.to(device)
saved_state_dict = torch.load(
RESTORE_FROM, map_location=lambda storage, loc: storage)
num_examples = saved_state_dict['example']
print("using running mean and running var")
log.writelines("using running mean and running var\n")
model.load_state_dict(saved_state_dict['state_dict'])
model.eval()
log.writelines('preprocessing time: ' +
str(time.time() - preprocess_start_time) + '\n')
print('Processing ' + slidename)
log.writelines('Processing ' + slidename + '\n')
TestTxt = os.path.join(DATA_LIST_PATH, slidename + '.txt')
testloader = data.DataLoader(LEAFTest(TestTxt,
crop_size=INPUT_SIZE,
mean=IMG_MEAN),
batch_size=BATCH_SIZE,
shuffle=False,
num_workers=16)
TestNpzPath = os.path.join(NPZ_PATH, slidename)
TestMapPath = os.path.join(MAP_PATH, slidename)
if not os.path.exists(TestNpzPath):
os.mkdir(TestNpzPath)
if not os.path.exists(TestMapPath):
os.mkdir(TestMapPath)
batch_time = AverageMeter()
with torch.no_grad():
end = time.time()
for index, (image, name) in enumerate(testloader):
output = model(image).to(device)
del image
Softmax = torch.nn.Softmax2d()
pred = torch.max(Softmax(output), dim=1, keepdim=True)
del output
for ind in range(0, pred[0].size(0)):
prob = torch.squeeze(pred[0][ind]).data.cpu().numpy()
prob = coo_matrix(prob)
if len(prob.data) == 0:
continue
mapname = name[ind].replace(
'.jpg', '_N' + str(num_examples) + '_MAP.npz')
mapfile = os.path.join(TestMapPath, mapname)
save_npz(mapfile, prob.tocsr())
msk = torch.squeeze(pred[1][ind]).data.cpu().numpy()
msk = coo_matrix(msk)
if len(msk.data) == 0:
continue
npzname = name[ind].replace(
'.jpg', '_N' + str(num_examples) + '_MSK.npz')
npzfile = os.path.join(TestNpzPath, npzname)
save_npz(npzfile, msk.tocsr())
batch_time.update(time.time() - end)
end = time.time()
if index % 10 == 0:
print(
'Test:[{0}/{1}]\t'
'Time {batch_time.val:.3f}({batch_time.avg:.3f})'.
format(index,
len(testloader),
batch_time=batch_time))
print('The total test time for ' + slidename + ' is ' +
str(batch_time.sum))
log.writelines('batch num:' + str(len(testloader)) + '\n')
log.writelines('The total test time for ' + slidename + ' is ' +
str(batch_time.sum) + '\n')
# merge npz
path = root_pth + tumorname + '/' + 'PatchNet/npz'
savenpz = os.path.join(path, 'whole_npz/' + model_id)
if not os.path.exists(savenpz):
os.makedirs(savenpz)
logdir = root_pth + tumorname + '/' + 'PatchNet/logfiles/'
if not os.path.exists(logdir):
os.makedirs(logdir)
logfile = open(
os.path.join(logdir, 'merge_npz_' + model_id + '.log'), 'w')
img = imread(test_img_pth)
width, height = img.shape[1], img.shape[0]
final = merge_npz(os.path.join(path, model_id), slidename,
os.path.join(savenpz, slidename + '_Map.npz'),
int(width), int(height))
# saving whole heatmap and mask
OldMatPath = root_pth + tumorname + '/' + 'PatchNet/npz/whole_npz/' + model_id
MatPath = root_pth + tumorname + '/' + 'PatchNet/npz/whole_npz_post/' + model_id
FigPath = root_pth + tumorname + '/' + 'PatchNet/npz/whole_fig_post/' + model_id
# CRFFigPath = "/data/AutoPheno/" + tumorname + "/PatchNet/npz/CRF"
PostPath = root_pth + tumorname + '/' + 'PatchNet/npz/whole_fig_post_rm/' + model_id
if not os.path.exists(MatPath):
os.makedirs(MatPath)
if not os.path.exists(FigPath):
os.makedirs(FigPath)
if not os.path.exists(PostPath):
os.makedirs(PostPath)
# if not os.path.exists(CRFFigPath):
# os.makedirs(CRFFigPath)
print("remove background")
rmbackground(test_img_pth,
OldMatPath,
MatPath,
ref_extent=ref_extent,
ref_area=ref_area)
SavePatchMap(MatPath, FigPath, slidename + "_Map.npz")
pred_img_pth = os.path.join(FigPath, slidename + "_Map.png")
# crf_img_name = slidename+"_Map.png"
# crf_img_pth = os.path.join(CRFFigPath, crf_img_name)
# CRFs(test_img_pth, pred_img_pth, crf_img_pth)
post_img_pth = os.path.join(PostPath, slidename + "_Map.png")
img = imread(pred_img_pth)
img = im2vl(img)
img_close = closing(img, square(3))
labeled_img = label(img_close, connectivity=2)
new_labeled_img = remove_small_objects(labeled_img,
min_size=rm_rf_area,
connectivity=2)
# remove non-circle region
props = regionprops(new_labeled_img)
for i, reg in enumerate(props):
if reg.eccentricity > ref_ecc:
new_labeled_img[new_labeled_img == reg.label] = 0
new_img = np.asarray(new_labeled_img != 0, dtype=np.uint8)
new_img = vl2im(new_img)
imsave(post_img_pth, new_img)
# Draw bounding box
cv2.namedWindow(imagelst[0], cv2.WINDOW_NORMAL)
imag = cv2.pyrDown(cv2.imread(post_img_pth, cv2.IMREAD_UNCHANGED))
# threshold image
ret, threshed_img = cv2.threshold(
cv2.cvtColor(imag, cv2.COLOR_BGR2GRAY), 127, 255,
cv2.THRESH_BINARY)
# find contours and get the external one
contours, hier = cv2.findContours(threshed_img, cv2.RETR_TREE,
cv2.CHAIN_APPROX_SIMPLE)
# with each contour, draw boundingRect in green
# a minEnclosingCircle in blue
rect_list = []
cir_list = []
for c in contours:
# get the bounding rect
x, y, w, h = cv2.boundingRect(c)
# draw a green rectangle to visualize the bounding rect
cv2.rectangle(imag, (x, y), (x + w, y + h), (0, 255, 0), 2)
rect_list.append([w * h, (x, y), (x + w, y + h)])
# finally, get the min enclosing circle
(x, y), radius = cv2.minEnclosingCircle(c)
# convert all values to int
center = (int(x), int(y))
radius = int(radius)
# and draw the circle in blue
imag = cv2.circle(imag, center, radius, (255, 0, 0), 2)
cir_list.append([radius, center])
postprocess(cir_list, rect_list, findN(cir_list), imag, imagelst)
cv2.drawContours(imag, contours, -1, (255, 255, 0), 1)
cv2.imshow(imagelst[0], imag)
del imagelst[0]
os.makedirs(PostPath)
# if not os.path.exists(CRFFigPath):
# os.makedirs(CRFFigPath)
print("remove background")
rmbackground(test_img_pth,
OldMatPath,
MatPath,
ref_extent=ref_extent,
ref_area=ref_area)
SavePatchMap(MatPath, FigPath, slidename + "_Map.npz")
pred_img_pth = os.path.join(FigPath, slidename + "_Map.png")
# crf_img_name = slidename+"_Map.png"
# crf_img_pth = os.path.join(CRFFigPath, crf_img_name)
# CRFs(test_img_pth, pred_img_pth, crf_img_pth)
post_img_pth = os.path.join(PostPath, slidename + "_Map.png")
img = imread(pred_img_pth)
img = im2vl(img)
img_close = closing(img, square(3))
labeled_img = label(img_close, connectivity=2)
new_labeled_img = remove_small_objects(labeled_img,
min_size=rm_rf_area,
connectivity=2)
# remove non-circle region
props = regionprops(new_labeled_img)
for i, reg in enumerate(props):
if reg.eccentricity > ref_ecc:
new_labeled_img[new_labeled_img == reg.label] = 0
new_img = np.asarray(new_labeled_img != 0, dtype=np.uint8)
new_img = vl2im(new_img)
imsave(post_img_pth, new_img)