"""
details
-------
1. Get name list of images and labels (/masks/ground truth)
2. Slice each case according to X/Y/Z dim
3. Save a list of [case number, slice number, slicename, labelname, average pixel value of the slice,
total pixel number of ground truth, bounding box (minA, maxA, minB, maxB)]
This code is adapted from https://github.com/198808xc/OrganSegC2F/blob/master/OrganSegC2F/init.py
"""
import numpy as np
import os
import sys
import time
from utils import *
# read input arguments
data_path = sys.argv[1] + "/"
organ_number = int(sys.argv[2])
folds = int(sys.argv[3])
low_range = int(sys.argv[4])
high_range = int(sys.argv[5])
if __name__=="__main__":
# get image list
image_list = []
image_filename = []
keyword = ''
for directory, _, file_ in os.walk(image_path):
for filename in sorted(file_):
if keyword in filename:
image_list.append(os.path.join(directory, filename))
image_filename.append(os.path.splitext(filename)[0])
# get label list
label_list = []
label_filename = []
for directory, _, file_ in os.walk(label_path):
for filename in sorted(file_):
if keyword in filename:
label_list.append(os.path.join(directory, filename))
label_filename.append(os.path.splitext(filename)[0])
# check if #image equals #labels
if len(image_list) != len(label_list):
exit('Error: the number of labels and the number of images are not equal!')
total_samples = len(image_list)
for plane in ['Z']:
output = open(list_training[plane], 'w')
output.close()
print 'Initialization starts.'
# iterate through all samples
for i in range(total_samples):
start_time = time.time()
print 'Processing ' + str(i + 1) + ' out of ' + str(len(image_list)) + ' files.'
image = np.load(image_list[i])
label = np.load(label_list[i])
# only z for now
for plane in ['Z']:
# slice_number is the number of slices of corresponding dimension (X/Y/Z)
slice_number = label.shape[2]
image_directory_ = os.path.join(image_path_[plane], image_filename[i])
if not os.path.exists(image_directory_):
os.makedirs(image_directory_)
label_directory_ = os.path.join(label_path_[plane], label_filename[i])
if not os.path.exists(label_directory_):
os.makedirs(label_directory_)
print ' Slicing data: ' + str(time.time() - start_time) + ' second(s) elapsed.'
# for storing the total number of pixels of ground truth mask
sum_ = np.zeros((slice_number, organ_number + 1), dtype = np.int)
# for storing bounding boxes of ground truth masks (A_min, A_max, B_min, B_max)
minA = np.zeros((slice_number, organ_number + 1), dtype = np.int)
maxA = np.zeros((slice_number, organ_number + 1), dtype = np.int)
minB = np.zeros((slice_number, organ_number + 1), dtype = np.int)
maxB = np.zeros((slice_number, organ_number + 1), dtype = np.int)
# for storing mean pixel value of each slice
average = np.zeros((slice_number), dtype = np.float)
# iterate through all slices of current case i and current plane
for j in range(0, slice_number):
# image_filename_ sample dir: image_X / 0001 / 0001.npy
# plane/ case num / slice num
image_filename_ = os.path.join( \
image_path_[plane], image_filename[i], '{:0>4}'.format(j) + '.npy')
label_filename_ = os.path.join( \
label_path_[plane], label_filename[i], '{:0>4}'.format(j) + '.npy')
image_ = image[:, :, j]
label_ = label[:, :, j]
# threshold image to specified range ([-100, 240] for pancreas)
image_[image_ < low_range] = low_range
image_[image_ > high_range] = high_range
# save sliced image and label
if not os.path.isfile(image_filename_) or not os.path.isfile(label_filename_):
np.save(image_filename_, image_)
np.save(label_filename_, label_)
# compute the mean value of the slice
average[j] = float(image_.sum()) / (image_.shape[0] * image_.shape[1])
for o in range(1, organ_number + 1):
# this is the sum of pixel numbers of a ground truth mask
sum_[j, o] = (is_organ(label_, o)).sum()
# record the coordinates of ground truth mask pixels
arr = np.nonzero(is_organ(label_, o))
# save the bounding box of ground truth mask (A_min, A_max, B_min, B_max)
minA[j, o] = 0 if not len(arr[0]) else min(arr[0])
maxA[j, o] = 0 if not len(arr[0]) else max(arr[0])
minB[j, o] = 0 if not len(arr[1]) else min(arr[1])
maxB[j, o] = 0 if not len(arr[1]) else max(arr[1])
# iterate each slice of current case i
for j in range(0, slice_number):
image_filename_ = os.path.join( \
image_path_[plane], image_filename[i], '{:0>4}'.format(j) + '.npy')
label_filename_ = os.path.join( \
label_path_[plane], label_filename[i], '{:0>4}'.format(j) + '.npy')
# append the following output to training_X/Y/Z.txt
output = open(list_training[plane], 'a+')
# case number, slice number
output.write(str(i) + ' ' + str(j))
# image file name, label file name
output.write(' ' + image_filename_ + ' ' + label_filename_)
# average pixel value of slice j, case i, and current plane
output.write(' ' + str(average[j]))
# sum of ground truth pixels, and bounding box of gt mask (A_min, A_max, B_min, B_max)
for o in range(1, organ_number + 1):
output.write(' ' + str(sum_[j, o]) + ' ' + str(minA[j, o]) + \
' ' + str(maxA[j, o]) + ' ' + str(minB[j, o]) + ' ' + str(maxB[j, o]))
output.write('\n')
output.close()
print ' ' + plane + ' plane is done: ' + \
str(time.time() - start_time) + ' second(s) elapsed.'
print 'Processed ' + str(i + 1) + ' out of ' + str(len(image_list)) + ' files: ' + \
str(time.time() - start_time) + ' second(s) elapsed.'
# create the 4 training image lists
print 'Writing training image list.'
for f in range(folds):
list_training_ = training_set_filename(f)
output = open(list_training_, 'w')
for i in range(total_samples):
if in_training_set(total_samples, i, folds, f):
output.write(str(i) + ' ' + image_list[i] + ' ' + label_list[i] + '\n')
output.close()
# create the 4 test image lists
print 'Writing testing image list.'
for f in range(folds):
list_testing_ = testing_set_filename(f)
output = open(list_testing_, 'w')
for i in range(total_samples):
if not in_training_set(total_samples, i, folds, f):
output.write(str(i) + ' ' + image_list[i] + ' ' + label_list[i] + '\n')
output.close()
print 'Initialization is done.'
