def get_model_data(model_name):
if model_name == 'Unet16':
model = unet.Unet(None, None, 16)
elif model_name == 'Unet32':
model = unet.Unet(None, None, 32)
elif model_name == 'Unet64':
model = unet.Unet(None, None, 64)
else:
os.sys.exit("bollcocks")
weights_fname = '{}.h5'.format(model_name)
submit_fname = '{}.csv'.format(model_name)
return model, weights_fname, submit_fname
def build_framework(self):
# Unet
self.output = unet.Unet(name="UNet",
in_data=self.input_data,
width=self.image_w,
height=self.image_h,
channel=self.image_c,
cal_num=self.class_num,
is_train=self.is_train,
reuse=False)
# loss
self.loss = unet.uu.generalized_dice_loss(self.output_mask,
self.output)
# optimizer
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
self.opt = tf.train.AdamOptimizer(learning_rate=self.lr).minimize(
self.loss, name="opt")
self.init = tf.global_variables_initializer()
# Add ops to save and restore all the variables.
self.saver = tf.train.Saver()
if self.GPU_str == 'yes':
self.sess = tf.Session(config=tf.ConfigProto(
log_device_placement=True))
else:
self.sess = tf.Session()
# whether training the model or load trained model
if self.tr_flag == 'yes':
self.sess.run(self.init)
else:
self.saver.restore(self.sess, self.model_path)
def __init__(self,lr,beta1,model_path,data_path,result_path,net_G_type="unet"):
self.device = torch.device('cuda:0')
self.train_dataset=mydataset.myDataset(data_path)
# self.test_dataset=mydataset.myDataset(data_path)
self.val_dataset=mydataset.myDataset(data_path,val_set=True)
self.result_path=result_path
self.model_path=model_path
self.data_path=data_path
if(net_G_type=="unet"):
self.net_G=unet.Unet(3,3,8).to(self.device)
else:
self.net_G=unet.GlobalGenerator(3,3,8).to(self.device)
#self.net_D=net_D.net_D(6).to(self.device)
self.net_D=net_D.MultiscaleDiscriminator(6, 64, 3, nn.BatchNorm2d, False, 1, True).to(self.device)
self.init_weights()
if os.path.exists(model_path+"/net_G.pth"):
self.net_G.load_state_dict(torch.load(model_path+"/net_G.pth"))
if os.path.exists(model_path+"/net_D.pth"):
self.net_D.load_state_dict(torch.load(model_path+"/net_D.pth"))
self.optimizer_G = torch.optim.Adam(self.net_G.parameters(), lr=lr, betas=(beta1, 0.999))
self.optimizer_D = torch.optim.Adam(self.net_D.parameters(), lr=lr, betas=(beta1, 0.999))
self.Tensor = torch.cuda.FloatTensor
self.criterionGAN = GANLoss(use_lsgan=True,tensor=self.Tensor).to(self.device)
self.criterionFeat = torch.nn.L1Loss().to(self.device)
self.real_label=torch.tensor(1.0)
self.fake_label=torch.tensor(0.0)
self.fake_pool = ImagePool(0)
self.criterionVGG = VGGLoss(0).to(self.device)
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
"dir",
help=
"Directory containing the image files (.png) we'll run inference on. \
Relative to the root of the project (tulip-fields/)")
args = parser.parse_args()
ctx = mx.gpu(0)
batch_size = 8
img_size = 256
root = os.path.dirname(__file__)
imgdir = os.path.join(root, os.pardir, args.dir)
checkpoint_dir = os.path.join(root, 'checkpoints', 'unet')
# Instantiate a U-Net and train it
net = unet.Unet()
net.load_params(os.path.join(checkpoint_dir, 'best_unet.params'), ctx)
print("Scanning dir {}".format(imgdir))
files = glob.glob(os.path.join(imgdir, '*wms*.png'))
print("Found {} images".format(len(files)))
nbatches = int(math.ceil(len(files) / batch_size))
reader = ImageReader(img_size, ctx)
for n in range(nbatches):
files_batch = files[n * batch_size:(n + 1) * batch_size]
batch = reader.load_batch(files_batch)
batch = batch.as_in_context(ctx)
preds = nd.argmax(net(batch), axis=1)
save_batch(files_batch, preds)
def __init__(self, lr, beta1, model_path, data_path, result_path):
self.device = torch.device('cuda:0')
self.train_dataset = mydataset.myDataset(data_path)
# self.test_dataset=mydataset.myDataset(data_path)
self.val_dataset = mydataset.myDataset(data_path, val_set=True)
self.result_path = result_path
self.model_path = model_path
self.data_path = data_path
self.net_G = unet.Unet(3, 3, 8).to(self.device)
self.net_D = net_D.net_D(6).to(self.device)
self.init_weights()
if os.path.exists(model_path + "/net_G.pth"):
self.net_G.load_state_dict(torch.load(model_path + "/net_G.pth"))
if os.path.exists(model_path + "/net_D.pth"):
self.net_D.load_state_dict(torch.load(model_path + "/net_D.pth"))
self.optimizer_G = torch.optim.Adam(self.net_G.parameters(),
lr=lr,
betas=(beta1, 0.999))
self.optimizer_D = torch.optim.Adam(self.net_D.parameters(),
lr=lr,
betas=(beta1, 0.999))
self.gan_loss = nn.MSELoss().to(self.device)
self.l1_loss = nn.L1Loss().to(self.device)
self.real_label = torch.tensor(1.0)
self.fake_label = torch.tensor(0.0)
def __init__(self, model, output):
super(UNetInferenceFn, self).__init__()
self.ctx = mx.cpu(0)
self.net = unet.Unet()
self.net.load_params(
os.path.join(os.path.dirname(model), 'best_unet.params'))
self.img_size = 256
self.reader = ImageReader(self.img_size, self.ctx)
self.output = output
def model_fn(features, labels, mode, params, config):
global_step = tf.train.get_or_create_global_step()
training = mode == tf.estimator.ModeKeys.TRAIN
net = unet.Unet(num_classes=params['data_loader'].num_classes)
logits = net(features['image'], training=training)
loss = losses.segmentation_loss(labels=labels['segmentation'],
logits=logits,
losses=params['losses'])
if mode == tf.estimator.ModeKeys.TRAIN:
optimizer = tf.train.AdamOptimizer(params['learning_rate'])
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_step = optimizer.minimize(loss, global_step=global_step)
build_summary(features['image'], labels['segmentation'],
logits) # TODO: refactor
summary_hook = tf.train.SummarySaverHook(
save_secs=60 * 3,
output_dir=config.model_dir,
summary_op=tf.summary.merge_all())
return tf.estimator.EstimatorSpec(mode,
loss=loss,
train_op=train_step,
training_hooks=[summary_hook])
elif mode == tf.estimator.ModeKeys.EVAL:
mask = tf.argmax(labels['segmentation'], -1)
predictions = tf.argmax(logits, -1)
metrics = {
'iou':
tf.metrics.mean_iou(labels=mask,
predictions=predictions,
num_classes=params['data_loader'].num_classes +
1)
}
build_summary(features['image'], labels['segmentation'], logits)
summary_hook = tf.train.SummarySaverHook(
save_steps=10,
# save_secs=60,
output_dir=os.path.join(config.model_dir, 'eval'),
summary_op=tf.summary.merge_all())
return tf.estimator.EstimatorSpec(mode,
loss=loss,
eval_metric_ops=metrics,
evaluation_hooks=[summary_hook])
def test_mask():
img = cv2.imread('sample_image.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
shape = img.shape
print(shape)
g = MaskGenerator((shape[0],shape[1]),43)
mask = g.random_mask()
masked_img = deepcopy(img)
masked_img[mask==0] = 255
_, axes = plt.subplots(1,3,figsize=(20,5))
axes[0].imshow(img)
axes[1].imshow(mask*255)
axes[2].imshow(masked_img)
import torch
import torchvision
totensor = torchvision.transforms.ToTensor()
formatted_img = masked_img
formatted_mask = np.expand_dims(mask,0)
formatted_img = totensor(formatted_img).unsqueeze(0)
formatted_mask = torch.from_numpy(formatted_mask.transpose(0, 3, 1, 2)).float()
print(torch.unique(formatted_img))
print(formatted_mask.size())
import unet
self = unet.Unet(3,64)
# self(torch.randn(1,3,512,512), torch.randn(1,3,512,512))
out_img1, out_mask1 = self.pconv1(formatted_img, formatted_mask)
out_img2, out_mask2 = self.pconv2(out_img1, out_mask1)
out_img3, out_mask3 = self.pconv3(out_img2, out_mask2)
out_img4, out_mask4 = self.pconv4(out_img3, out_mask3)
out_img5, out_mask5 = self.pconv5(out_img4, out_mask4)
out_img6, out_mask6 = self.pconv6(out_img5, out_mask5)
out_img7, out_mask7 = self.pconv7(out_img6, out_mask6)
out_img8, out_mask8 = self.pconv8(out_img7, out_mask7)
#
_,axes = plt.subplots(2,4,figsize=(20,5))
axes[0][0].imshow(out_mask1[0,0,:,:],cmap='gray',vmin=0,vmax=1)
axes[0][1].imshow(out_mask2[0, 0, :, :], cmap='gray', vmin=0, vmax=1)
axes[0][2].imshow(out_mask3[0, 0, :, :], cmap='gray', vmin=0, vmax=1)
axes[0][3].imshow(out_mask4[0, 0, :, :], cmap='gray', vmin=0, vmax=1)
axes[1][0].imshow(out_mask5[0, 0, :, :], cmap='gray', vmin=0, vmax=1)
axes[1][1].imshow(out_mask6[0, 0, :, :], cmap='gray', vmin=0, vmax=1)
axes[1][2].imshow(out_mask7[0, 0, :, :], cmap='gray', vmin=0, vmax=1)
axes[1][3].imshow(out_mask8[0, 0, :, :], cmap='gray', vmin=0, vmax=1)
plt.show()
def main(args):
# load correct model
if args.attention == True:
model = unet.AttentionUnet(in_channels=NUM_CHANNELS,
out_channels=NUM_CLASSES)
elif args.attention == False:
# we don't use residual for non-denoising purposes
model = unet.Unet(in_channels=NUM_CHANNELS, out_channels=NUM_CLASSES)
# check if we use GPU
if args.cuda:
model = model.cuda()
# check for mode - training or evaluation?
if args.mode == 'eval':
evaluate(args, model)
elif args.mode == 'train':
train(args, model)
def model_fn(model_dir):
"""function used to load pretrained model"""
ctx = mx.cpu()
net = unet.Unet()
print("Loading", model_dir)
if path.exists(model_dir + "/unet_RGB.params"):
print("Loading RGB Model")
net.load_params(model_dir + "/unet_RGB.params", ctx)
print("RGB Model Loaded")
elif path.exists(model_dir + "/unet_ALL_BANDS.params"):
print("Loading ALL_BANDS Model")
net.load_params(model_dir + "/unet_ALL_BANDS.params", ctx)
print("ALL_BANDS Model Loaded")
else:
print("Model Missing")
net = None
return (net)
def predict_model(name):
para = Parameter()
root_address = para.root_address
data_address = os.path.join(root_address, 'data/test')
prediction_save_address = os.path.join(root_address, 'result/pre_image')
unet_model_path = os.path.join(
root_address, 'result/unet_trained/model.ckpt/{}'.format(name))
provider_path = os.path.join(root_address, 'data/test/*/*/*.npy')
if not os.path.exists(prediction_save_address):
os.mkdir(prediction_save_address)
patient_pre_list = []
patient_label_list = []
patient_save_list = []
for addr in os.listdir(data_address): #addr:01
image_pre_list = []
image_label_list = []
image_save_list = []
patient_addr = os.path.join(data_address, addr) #data/test/01
patient_save_addr = os.path.join(prediction_save_address,
addr) #pre_image/01
if not os.path.exists(patient_save_addr):
os.mkdir(patient_save_addr)
for addr_1 in os.listdir(patient_addr): #addr_1:0
image_addr = os.path.join(patient_addr, addr_1) #data/test/01/0
images_save_addr = os.path.join(patient_save_addr,
addr_1) #pre_image/01/0
if not os.path.exists(images_save_addr):
os.mkdir(images_save_addr)
for filename in os.listdir(image_addr):
if 'fat.npy' in filename:
image_pre_list.append(os.path.join(
image_addr, filename)) #data/test/01/0/.._fat.npy
image_save_list.append(
os.path.join(images_save_addr, filename).replace(
'.npy', '_pre.npy')) #pre_image01/0/.._fat_pre.npy
if 'label.npy' in filename:
image_label_list.append(os.path.join(
image_addr, filename)) #data/test/01/0/.._label.npy
patient_pre_list.append(image_pre_list)
patient_label_list.append(image_label_list)
patient_save_list.append(image_save_list)
print('start to predict')
if para.channel == 1:
generator = image_gen.oneChannelProvider(provider_path)
elif para.channel == 4:
generator = image_gen.fourChannelProvider(provider_path)
elif para.channel == 8:
generator = image_gen.eightChannelProvider(provider_path)
net = unet.Unet(channels=generator.channels,
n_class=generator.n_class,
cost=para.cost,
cost_kwargs=dict(regularizer=para.regularizer),
layers=para.layers,
features_root=para.features_root,
training=False)
init = tf.global_variables_initializer()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
# Initialize variables
sess.run(init)
# Restore model weights from previously saved model
net.restore(sess, unet_model_path)
for i in range(len(patient_pre_list)):
print('predict on {}'.format(patient_pre_list[i][0]))
for j in range(len(patient_pre_list[i])):
if generator.channels == 1:
path = patient_pre_list[i][j]
fat_path = path.replace("fat", "inn")
fat_img = np.array(np.load(fat_path), dtype=np.float32)
img_arr = np.zeros((fat_img.shape[0], fat_img.shape[1], 1),
dtype=np.float32)
img_arr[..., 0] = fat_img
x_test = img_arr.reshape(1, img_arr.shape[0],
img_arr.shape[1],
generator.channels)
y_dummy = np.empty(
(x_test.shape[0], x_test.shape[1], x_test.shape[2], 2))
prediction = sess.run(net.predicter,
feed_dict={
net.x: x_test,
net.y: y_dummy,
net.keep_prob: 1.
})
mask = prediction[0, ..., 1] > para.mask
#mask = prediction[0,...,1]
np.save(patient_save_list[i][j], mask)
plt.imshow(mask, cmap='gray')
plt.savefig(patient_save_list[i][j].replace(
'.npy', '.png'))
elif generator.channels == 4:
path = patient_pre_list[i][j]
fat_path = path.replace("fat", "fat")
inn_path = path.replace("fat", "inn")
wat_path = path.replace("fat", "wat")
opp_path = path.replace("fat", "opp")
fat_img = np.array(np.load(fat_path), dtype=np.float32)
inn_img = np.array(np.load(inn_path), dtype=np.float32)
wat_img = np.array(np.load(wat_path), dtype=np.float32)
opp_img = np.array(np.load(opp_path), dtype=np.float32)
img_arr = np.zeros((fat_img.shape[0], fat_img.shape[1], 4),
dtype=np.float32)
img_arr[..., 0] = fat_img
img_arr[..., 1] = inn_img
img_arr[..., 2] = wat_img
img_arr[..., 3] = opp_img
x_test = img_arr.reshape(1, img_arr.shape[0],
img_arr.shape[1],
generator.channels)
y_dummy = np.empty(
(x_test.shape[0], x_test.shape[1], x_test.shape[2], 2))
prediction = sess.run(net.predicter,
feed_dict={
net.x: x_test,
net.y: y_dummy,
net.keep_prob: 1.
})
mask = prediction[0, ..., 1] > para.mask
#mask = prediction[0,...,1]
np.save(patient_save_list[i][j], mask)
plt.imshow(mask, cmap='gray')
plt.savefig(patient_save_list[i][j].replace(
'.npy', '.png'))
elif generator.channels == 8:
path = patient_pre_list[i][j]
fat_path = path.replace("fat", "fat")
inn_path = path.replace("fat", "inn")
wat_path = path.replace("fat", "wat")
opp_path = path.replace("fat", "opp")
fin_path = path.replace("fat", "fin")
win_path = path.replace("fat", "win")
wop_path = path.replace("fat", "wop")
iop_path = path.replace("fat", "iop")
fat_img = np.array(np.load(fat_path), dtype=np.float32)
inn_img = np.array(np.load(inn_path), dtype=np.float32)
wat_img = np.array(np.load(wat_path), dtype=np.float32)
opp_img = np.array(np.load(opp_path), dtype=np.float32)
fin_img = np.array(np.load(fin_path), dtype=np.float32)
win_img = np.array(np.load(win_path), dtype=np.float32)
wop_img = np.array(np.load(wop_path), dtype=np.float32)
iop_img = np.array(np.load(iop_path), dtype=np.float32)
img_arr = np.zeros((fat_img.shape[0], fat_img.shape[1], 8),
dtype=np.float32)
img_arr[..., 0] = fat_img
img_arr[..., 1] = inn_img
img_arr[..., 2] = wat_img
img_arr[..., 3] = opp_img
img_arr[..., 4] = fin_img
img_arr[..., 5] = win_img
img_arr[..., 6] = wop_img
img_arr[..., 7] = iop_img
x_test = img_arr.reshape(1, img_arr.shape[0],
img_arr.shape[1],
generator.channels)
y_dummy = np.empty(
(x_test.shape[0], x_test.shape[1], x_test.shape[2], 2))
prediction = sess.run(net.predicter,
feed_dict={
net.x: x_test,
net.y: y_dummy,
net.keep_prob: 1.
})
mask = prediction[0, ..., 1] > para.mask
#mask = prediction[0,...,1]
np.save(patient_save_list[i][j], mask)
plt.imshow(mask, cmap='gray')
plt.savefig(patient_save_list[i][j].replace(
'.npy', '.png'))
elif generator.channels == 12:
path = patient_pre_list[i][j]
fat_path = path.replace("fat", "fat")
inn_path = path.replace("fat", "inn")
wat_path = path.replace("fat", "wat")
opp_path = path.replace("fat", "opp")
fat_img = np.array(np.load(fat_path), dtype=np.float32)
inn_img = np.array(np.load(inn_path), dtype=np.float32)
wat_img = np.array(np.load(wat_path), dtype=np.float32)
opp_img = np.array(np.load(opp_path), dtype=np.float32)
img_arr = np.zeros(
(fat_img.shape[1], fat_img.shape[2], 12),
dtype=np.float32)
img_arr[..., 0] = fat_img[0]
img_arr[..., 1] = fat_img[1]
img_arr[..., 2] = fat_img[2]
img_arr[..., 3] = inn_img[0]
img_arr[..., 4] = inn_img[1]
img_arr[..., 5] = inn_img[2]
img_arr[..., 6] = wat_img[0]
img_arr[..., 7] = wat_img[1]
img_arr[..., 8] = wat_img[2]
img_arr[..., 9] = opp_img[0]
img_arr[..., 10] = opp_img[1]
img_arr[..., 11] = opp_img[2]
x_test = img_arr.reshape(1, img_arr.shape[0],
img_arr.shape[1],
generator.channels)
y_dummy = np.empty(
(x_test.shape[0], x_test.shape[1], x_test.shape[2], 2))
prediction = sess.run(net.predicter,
feed_dict={
net.x: x_test,
net.y: y_dummy,
net.keep_prob: 1.
})
mask = prediction[0, ..., 1] > para.mask
np.save(patient_save_list[i][j], mask)
plt.imshow(mask, cmap='gray')
plt.savefig(patient_save_list[i][j].replace(
'.npy', '.png'))
name = "myo_cv_" + str(i + 1)
# make correct data sets
train_ind = np.concatenate(np.delete(folds_ind, i, 0))
test_ind = folds_ind[i]
print(train_ind)
print(test_ind)
# test_ind = test_inds[i]
#load coresponding 5 fold cv validation set
image_test_ds = image_ds[test_ind, :, :, :]
masks_test_ds = masks_ds[test_ind, :, :, :]
vent_val_params = params.Params(dim=dim, experiment_name=name)
model = unet.Unet(vent_val_params)
predicted_masks_cv = model.make_prediction(image_test_ds, thresh=0.5)
print(image_test_ds.shape)
print(predicted_masks_cv.shape)
for j in range(0, image_test_ds.shape[0]):
dice[i, j] = unet.dice_coeff_np(masks_test_ds[j, :, :, 0],
predicted_masks_cv[j, :, :, 0])
# ROC
predicted_masks_roc = model.make_prediction(image_test_ds, thresh=None)
fpr, tpr, thresholds = roc_curve(masks_test_ds.flatten(),
predicted_masks_roc.flatten())
plt.plot(fpr,
val_loader = DataLoader(val_data,
batch_size=1,
shuffle=False,
num_workers=args.workers)
down = [(5, 5, 5), (5, 5, 5), (5, 5, 5), (5, 5, 5)]
up = [(5, 5, 5), (5, 5, 5), (5, 5, 5)]
if args.model in ('harmonic', 'unconstrained'):
network = hunet.HUnet(in_features=3, down=down, up=up, radius=2)
elif args.model == 'baseline':
down = [unet.repr_to_n(d) for d in down]
up = [unet.repr_to_n(d) for d in up]
network = unet.Unet(up=up, down=down, in_features=3)
cuda = torch.cuda.is_available()
network_repr = repr(network)
print(network_repr)
writer.add_text('general', network_repr)
if args.model == 'unconstrained':
for module in network.modules():
if hasattr(module, 'relax'):
module.relax()
print(f'relaxing {repr(module)}')
network_repr = repr(network)
print(network_repr)
def main():
parser = argparse.ArgumentParser(description="Train U-net")
parser.add_argument('--name', type=str, default='unknown',
help='network name')
parser.add_argument('--model_dir', type=str, required=True,
help='Where network will be saved and restored')
parser.add_argument("--lr",
type=float,
default=1e-4,
help="Adam learning rate")
parser.add_argument("--batch_size",
type=int,
default=5,
help="Batch size")
parser.add_argument("--input_size",
type=int,
default=324,
help="Input size of the image will fed into network. Input_size = 16*n + 4, Default: 324")
parser.add_argument("--output_size",
type=int,
default=116,
help="size of the image produced by network. Default: 116")
parser.add_argument("--tb_log_dir",
type=str,
required=True,
help="Tensorboard log dir")
parser.add_argument("--n_steps",
type=int,
default=0,
help="Number of the steps. Default: 0 means infinity steps.")
parser.add_argument("--dataset_dir",
type=str,
default="../dataset/trainset")
parser.add_argument("--pretrained_vgg",
type=str,
choices=['yes', 'no'],
default="yes",
help="Use pretrained vgg weigth")
parser.add_argument("--fix_vgg",
type=str,
choices=['yes', 'no'],
default="yes",
help="Fix vgg weights while learning")
parser.add_argument("--validation_freq",
type=int,
default=100,
help="Validation freq. Default 100")
parser.add_argument("--validation_set_size",
type=int,
default=20,
help="metrics will be averaged by validation_set_size. Default 20")
parser.add_argument("--channel",
type=str,
choices=['rgb', 'gray'],
default="rgb",
help="channel. Default: rgb")
args = parser.parse_args()
net_name = args.name
model_dir = args.model_dir
learning_rate = args.lr
batch_size = args.batch_size
net_input_size = args.input_size
net_output_size = args.output_size
tb_log_dir = args.tb_log_dir
n_steps = args.n_steps
dataset_dir = args.dataset_dir
pretrained_vgg = args.pretrained_vgg == 'yes'
fix_vgg = args.fix_vgg == 'yes'
validation_freq = args.validation_freq
validation_set_size = args.validation_set_size
channel = args.channel
print("Load dataset")
dataset = DS.DataSet(dataset_dir)
print("Initialize network manager")
network_manager = NManager(model_dir, net_name)
if network_manager.registered:
net = network_manager.get_net()
else:
print("Use pretrained weihts %s" % str(pretrained_vgg))
net = U.Unet(vgg_pretrained=pretrained_vgg)
network_manager.register_net(net)
print("Move to GPU")
net.cuda()
if channel == "rgb":
def get_features(x):
return x.get_ndarray([DS.ChannelRGB_PanSharpen])
else:
def get_features(x):
img0 = x.get_ndarray([DS.ChannelPAN])[0]
img = np.array([img0, img0, img0])
return img
def get_target(x):
return x.get_interior_mask()
train_sampler = S.Sampler(dataset.train_images(), get_features, get_target,
net_input_size, net_output_size, rotate_amplitude=20,
random_crop=True, reflect=True)()
test_sampler = S.Sampler(dataset.test_images(), get_features, get_target,
net_input_size, net_output_size, rotate_amplitude=20,
random_crop=True, reflect=True)()
if fix_vgg:
parameters = list(net.bn.parameters()) + list(net.decoder.parameters()) + list(net.conv1x1.parameters())
else:
parameters = net.parameters()
print("LR: %f" % learning_rate)
optimizer = torch.optim.Adam(parameters, lr=learning_rate)
logger = SummaryWriter(tb_log_dir + "/" + net_name)
print("Start learning")
with network_manager.session(n_steps) as (iterator, initial_step):
for step in tqdm.tqdm(iterator, initial=initial_step):
batch_features, batch_target = batch_generator(train_sampler, batch_size)
batch_features = Variable(FloatTensor(batch_features)).cuda()
batch_target = Variable(FloatTensor(batch_target)).cuda()
predicted = net.forward(batch_features)
train_metrics = eval_base_metrics(predicted, batch_target)
train_metrics = eval_precision_recall_f1(**train_metrics)
loss = train_metrics['loss']
optimizer.zero_grad()
loss.backward()
optimizer.step()
log_metrics(logger, '', train_metrics, step)
logger.add_scalar('lr', np.log(learning_rate)/np.log(10), step)
if step % 1000 == 0:
network_manager.save()
if step % validation_freq == 0:
test_metrics = average_metrics(net, test_sampler, batch_size, validation_set_size)
log_metrics(logger, 'val', test_metrics, step)
avg_train_metrics = average_metrics(net, train_sampler, batch_size, validation_set_size)
log_metrics(logger, 'avg_train', avg_train_metrics, step)
generate_image(logger, net, 'val', dataset.test_images(), get_features, get_target,
net_input_size, net_output_size, step)
generate_image(logger, net, 'train', dataset.train_images(), get_features, get_target,
net_input_size, net_output_size, step)
def main():
parser = argparse.ArgumentParser(
description='Script to segment tulip fields from the input images')
parser.add_argument(
"params",
help="Path to .params file with the saved UNet parameters",
metavar='net_params')
parser.add_argument(
"dirs",
help="Directories containing the image files we'll run inference on. \
Separated by whitespaces and relative to the root of the project (tulip-fields/)",
metavar="dirs",
nargs="*")
parser.add_argument(
"--output",
help="Path to directory where predicted tulip masks are saved. \
If not specified, predictions are saved in the same folder as their corresponding image",
dest="out",
metavar="output_dir",
default=None)
parser.add_argument(
"--mask-dir",
help="Path to directory containing the ground truth masks. \
If not specified, output plots will only contain predictions",
dest="masks",
default=None)
parser.add_argument("--batch-size",
help="Batch size (default is 8)",
dest="batch_size",
default=8,
type=int)
parser.add_argument(
"--ext",
help="Extension of the image files to be read (default is .png)",
dest="ext",
default=".png")
parser.add_argument(
"--multispectral",
help="Enables 13 band multispectral image mode (default is only RGB)",
dest="multisp",
action="store_true")
parser.add_argument("--use-cpu",
help="Use CPU instead of GPU",
dest="cpu",
action="store_true")
parser.add_argument(
"--show",
help="If set, the predictions for the first N batches are displayed",
default=0,
metavar="N",
type=int)
args = parser.parse_args()
ctx = mx.cpu() if args.cpu else mx.gpu(0)
root = os.path.join(os.path.dirname(__file__), os.pardir)
img_dirs = [os.path.join(root, d) for d in args.dirs]
outdir = os.path.join(root, args.out) if args.out else None
# Instantiate a U-Net and load the saved state
net = unet.Unet()
net.load_params(os.path.join(root, args.params), ctx)
loader = BatchLoader(img_dirs, args.ext, args.batch_size, ctx,
args.multisp)
if len(loader) == 0:
logging.error('No images found, please specify another directory')
else:
for idx, (batch, filenames) in enumerate(loader.get_batches()):
preds = nd.argmax(net(batch), axis=1)
save_batch(filenames, preds, outdir)
if idx < args.show:
plot_predictions(batch, preds, filenames, args.masks,
args.multisp)
# Subgrid stress
tau = [dataset['t' + c].data for c in comps]
# Deviatoric subgrid stress
tr_tau = tau[0] + tau[1]
tau[0] = tau[0] - tr_tau / 2
tau[1] = tau[1] - tr_tau / 2
# Reshape as (batch, *shape, channels)
inputs = np.moveaxis(np.array(S), 0, -1)[None]
labels = np.moveaxis(np.array(tau), 0, -1)[None]
return inputs, labels
# Build network and optimizer
tf.random.set_seed(tf_seed)
model = unet.Unet(stacks, stack_width, filters_base, output_channels,
**unet_kw)
optimizer = tf.keras.optimizers.Adam(learning_rate)
checkpoint = tf.train.Checkpoint(optimizer=optimizer, net=model)
if restore_epoch:
restore_path = f"{checkpoint_path}-{restore_epoch}"
checkpoint.restore(restore_path) #.assert_consumed()
print('Restored from {}'.format(restore_path))
else:
print('Initializing from scratch.')
initial_epoch = checkpoint.save_counter.numpy() + 1
def array_of_tf_components(tf_tens):
"""Create object array of tensorflow packed tensor components."""
# Collect components
# Tensorflow shaped as (batch, *shape, channels)
U=U,
sm=sm,
lm=lm,
Um=Um,
cn=cn,
dataset=dataset)
validloader = data.DataLoader(loader,
batch_size=batch,
num_workers=0,
shuffle=False,
drop_last=False)
hard_negative = 0
if net == 'U':
model = unet.Unet(feature_scale=8, input_size=input_layers)
if net == 'P':
model = pixel_net.PixelNet(K=net_K, input_size=input_layers)
model = model.cuda()
optimizer = optim.Adam(model.parameters(),
lr=init_lr,
betas=(0.9, 0.999),
eps=1e-8,
weight_decay=1e-5)
display_losses = drawloss.DisplayLosses(evaluate_it=5)
itt = -1
while itt < iterace:
plt.show()
#instantiate params
name = "cone"
batch_size = 1
num_epochs = 130
steps_per = 50
num_folds = 0
dim = image_ds.shape[1:3]
cone_train_params = params.Params(dim=dim,
experiment_name=name,
batch_size=batch_size,
num_epochs=num_epochs,
steps_per_epoch=steps_per)
image_train_ds, image_test_ds, masks_train_ds, masks_test_ds = train_test_split(
image_ds, masks_ds, test_size=0.1)
#make dynamic
model = unet.Unet(cone_train_params)
model.get_Unet()
print(model.params.callbacks)
model.train(image_train_ds,
masks_train_ds,
image_test_ds,
masks_test_ds,
continue_train=True)
#model.train(image_ds,masks_ds,continue_train = True)
def main():
parser = argparse.ArgumentParser(description="Script to train a U-Net on the images passed as parameters.")
parser.add_argument("--train", metavar="train_images",
help="Path to directory containing the training images (relative to project root)",
default="/opt/ml/input/data/train")
parser.add_argument("--val", metavar="val_images",
help="Path to directory containing the validation images",
default="/opt/ml/input/data/validation")
parser.add_argument("--masks", metavar="masks",
help="Path to directory containing the ground truth masks",
default="/opt/ml/input/data/masks")
parser.add_argument("--train-masks",
help="Used to define a different directory for training set masks",
default=None, dest="train_masks")
parser.add_argument("--checkpoints", metavar="checkpoints_dir",
help="Path to directory where the parameters of the best models will be stored",
default="/opt/ml/model/")
parser.add_argument("--epochs",
help="Number of epochs to train for (default is 50)",
type=int, default=50)
parser.add_argument("--batch-size",
help="Specify batch size (default is 8)",
type=int, default=8)
parser.add_argument("--multispectral",
help="Enable multispectral mode",
default="rgb")
parser.add_argument("--use-cpu",
help="Train on CPU instead of GPU (very slow, not recommended)",
dest="cpu", action="store_true")
parser.add_argument("--gpu-count",
help="Number of GPUs to use",
type=int, default=1, dest="gpus")
parser.add_argument("--use-eia",
help="Train on Elastic Inference",
dest="eia", action="store_true")
parser.add_argument("--drop-bands",
help="Indexes of bands not to be used, separed by commas",
default=None, dest="drop")
parser.add_argument("--show",
help="Plots the learning curve once training is completed",
dest="show", action="store_true")
parser.add_argument("--data-aug",
help="If selected, training data is augmented online before being fed to the network",
dest="aug", action="store_true")
args = parser.parse_args()
if args.eia:
ctx = [mx.eia()]
elif args.cpu:
ctx = [mx.cpu()]
else:
ctx=[mx.gpu(i) for i in range(args.gpus)]
if args.multispectral == "rgb":
multispectral=False
else:
multispectral=True
batch_size = args.batch_size
root = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))
train_dir = os.path.join(root, args.train)
val_dir = os.path.join(root, args.val)
mask_dir = os.path.join(root, args.masks)
checkpoint_dir = os.path.join(root, args.checkpoints)
tmp = os.path.join(os.path.dirname(__file__), 'tmp')
train_masks = os.path.join(root, args.train_masks) if args.train_masks is not None else mask_dir
drop_bands = [int(band) for band in args.drop.split(',')] if args.drop is not None else None
# Create train and validation DataLoaders from our Datasets
if args.aug:
print ("Using AugmentorDataset")
dataset_class = AugmentorDataset
else:
dataset_class = ImageWithMaskDataset
train_ds = dataset_class(train_dir, train_masks, multisp=multispectral,
transform_fn=lambda b, m, ms: transform(b, m, ms, drop=drop_bands))
train_iter = gluon.data.DataLoader(train_ds, batch_size, shuffle=True)
val_ds = ImageWithMaskDataset(val_dir, mask_dir, multisp=multispectral,
transform_fn=lambda b, m, ms: transform(b, m, ms, drop=drop_bands))
val_iter = gluon.data.DataLoader(val_ds, batch_size, shuffle=True)
if len(train_ds) == 0:
raise ValueError('The train directory {} does not contain any valid images'.format(train_dir))
if len(val_ds) == 0:
raise ValueError('The test directory {} does not contain any valid images'.format(train_dir))
data = {'train': train_iter, 'val': val_iter}
# Instantiate a U-Net and train it
net = unet.Unet()
net.collect_params().initialize(mx.init.Xavier(), ctx=ctx)
net.hybridize()
loss = DiceCoeffLoss()
trainer = gluon.Trainer(net.collect_params(), 'adam',
{'learning_rate': 1e-4, 'beta1': 0.9, 'beta2': 0.99})
epochs = args.epochs
results = train(net, data, loss, trainer, ctx, epochs, tmp)
if args.show:
import matplotlib.pyplot as plt
plt.plot(range(len(results['val'])), results['val'])
plt.title('Model learning curve')
plt.xlabel('Epoch')
plt.ylabel('IoU score')
plt.show()
# Find best scoring model
best = results['val'].index(max(results['val']))
best_params = os.path.join(tmp, results['names'][best])
# Copy it to <checkpoints> folder
os.makedirs(checkpoint_dir, exist_ok=True)
if multispectral:
bands = 'ALL_BANDS' if drop_bands is None else ('ALL_BANDS-' + '-'.join([str(b) for b in drop_bands]))
else:
bands = 'RGB'
save_filename = os.path.join(checkpoint_dir, 'unet_{}.params'.format(bands))
shutil.copyfile(best_params, save_filename)
np.save(os.path.join(checkpoint_dir, 'unet_{}_learning_curve_data.npy'.format(bands)), results)
shutil.rmtree(tmp)
print ('Best model on validation set saved in: {}'.format(save_filename))
height = 160
width = 160
batch_size = 10
train_path = '/DB/rhome/qyzheng/Desktop/qyzheng/source/renji_data/from_senior/0_cv_train.csv'
val_path = '/DB/rhome/qyzheng/Desktop/qyzheng/source/renji_data/from_senior/0_cv_val.csv'
dataset, iters = image_gen.GetDataset(train_path, batch_size)
generator = image_gen.BladderDataProvider(height, width, dataset)
"""
x_test, y_test = generator(1)
fig, ax = plt.subplots(1, 2, sharey=True, figsize=(8, 4))
ax[0].imshow(x_test[0, ..., 0], aspect="auto")
ax[1].imshow(y_test[0, ..., 1], aspect="auto")
#plt.show()
"""
net = unet.Unet(channels=generator.channels, n_class=generator.n_class, layers=4, features_root=64)
trainer = unet.Trainer(net, batch_size=4, optimizer="momentum", opt_kwargs=dict(momentum=0.2))
path = trainer.train(generator, "../unet_trained", training_iters=iters, epochs=100, display_step=4,
prediction_path='/DATA/data/sxfeng/data/IVDM3Seg/result/result_2/prediction')
'''
x_test, y_test = generator(1)
print(x_test.shape)
print(y_test.shape)
prediction = net.predict("../unet_trained/model.ckpt", x_test)
print(prediction.shape)
'''
"""
fig, ax = plt.subplots(1, 3, sharex=True, sharey=True, figsize=(12, 5))
ax[0].imshow(x_test[0,...,0], aspect="auto")
ax[1].imshow(y_test[0,...,1], aspect="auto")
cv_scores = []
fold_ind = np.arange(0, image_ds.shape[0])
np.random.shuffle(fold_ind)
folds_ind = np.array_split(fold_ind, num_folds)
for i in range(0, num_folds):
print("fold: " + str(i + 1))
name = "myo_cv_" + str(i + 1)
myo_train_params = params.Params(dim=dim,
experiment_name=name,
batch_size=batch_size,
num_epochs=num_epochs,
steps_per_epoch=steps_per)
model = unet.Unet(myo_train_params)
train_ind = np.concatenate(np.delete(folds_ind, i, 0))
test_ind = folds_ind[i]
print(train_ind)
print(test_ind)
image_train_ds = image_ds[train_ind, :, :, :]
masks_train_ds = masks_ds[train_ind, :, :, :]
image_test_ds = image_ds[test_ind, :, :, :]
masks_test_ds = masks_ds[test_ind, :, :, :]
#for i in range(0,image_train_ds.shape[0]):
# plt.subplot(1,2,1)
# plt.imshow(image_train_ds[i,:,:,0])
# plt.subplot(1,2,2)
# plt.imshow(masks_train_ds[i,:,:,0])
load_cone=False) #save numpys in relavent dirs np.save(numpy_path + "image_myo_ds", image_myo_ds) np.save(numpy_path + "masks_myo_ds", masks_myo_ds) image_myo_ds = np.load(numpy_path + "image_myo_ds.npy") masks_myo_ds = np.load(numpy_path + "masks_myo_ds.npy") print(image_myo_ds.shape) print(masks_myo_ds.shape) size = [800, 800] predict_cone_params = params.Params(dim=size, experiment_name="cone") model = unet.Unet(predict_cone_params) predicted_masks = model.make_prediction(image_myo_ds) image_nc_ds = np.multiply(image_myo_ds, predicted_masks) np.save(numpy_path + "image_nc_ds", image_nc_ds) # pad_mask[pad_mask == 1] = 0 # vent mask # if(np.amax(pad_mask[:,:,0] == 2.0)): # only if there are multiple masks # pad_mask[pad_mask == 1] = 0 # decide to seg cone instead of ventricle # max_val = np.amax(pad_mask[:,:,0]) # image = f._process_pathname(image_path_2CH + img) # X, Y, Z, T = image.shape
fold_ind = np.arange(0,image_ds.shape[0])
np.random.shuffle(fold_ind)
folds_ind = np.array_split(fold_ind,num_folds)
#for i in range (0,len(folds_ind)):
i = 1
print("fold: " + str(i+1))
name = "vent_cv_" + str(i+1)
vent_train_params = params.Params(dim = dim,
experiment_name = name,
batch_size = batch_size,
num_epochs = num_epochs,
steps_per_epoch = steps_per)
model = unet.Unet(vent_train_params)
train_ind = np.concatenate(np.delete(folds_ind,i))
test_ind = folds_ind[i]
image_train_ds = image_ds[train_ind,:,:,:]
masks_train_ds = masks_ds[train_ind,:,:,:]
image_test_ds = image_ds[test_ind,:,:,:]
masks_test_ds = image_ds[test_ind,:,:,:]
model.train(image_train_ds,masks_train_ds,image_test_ds,masks_test_ds)
scores = model.evaluate(image_test, masks_test, verbose = 1)
cv_scores.append(scores)
print(scores)
#print("%.2f% (+/- %.2f%)" % (np.mean(cv_scores), np.std(cv_scores)))
network = hunet.HUnet(in_features=4,
down=down,
up=up,
radius=2,
setup=setup)
elif args.model == 'baseline':
if args.dropout is not None:
dropout = functools.partial(torch.nn.Dropout2d, p=args.dropout)
setup = {**unet.default_setup, 'dropout': dropout}
else:
setup = unet.default_setup
down = [unet.repr_to_n(d) for d in down]
up = [unet.repr_to_n(d) for d in up]
network = unet.Unet(up=up, down=down, in_features=4, setup=setup)
cuda = torch.cuda.is_available()
network_repr = repr(network)
print(network_repr)
writer.add_text('general', network_repr)
n_params = 0
for param in network.parameters():
n_params += param.numel()
print(n_params, 'learnable parameters')
if cuda:
network = network.cuda()
import params
import unet
path = "/home/zeke/Programming/cnn/us_seg/"
load_path = "data/whole_dataset/4CH/"
masks_save_path = "data/4CH_dataset/masks/"
image_save_path = "data/4CH_dataset/image/"
#instantiate params
dim = [800, 800]
data_generation_params_vent = params.Params(dim=dim,
experiment_name="2CH_vent")
data_generation_params_cone = params.Params(dim=dim, experiment_name="cone")
#Load model and weights
model_vent = unet.Unet(data_generation_params_vent)
model_cone = unet.Unet(data_generation_params_cone)
for filename in os.listdir(load_path):
#Open image with Iksnap
subprocess.call(['itksnap', load_path + filename])
#Select Frames
frames = []
frame = 1
while (frame != 0):
frame = int(input("Enter frame(0 to quit): "))
if (frame != 0):
frames.append(frame)
if (len(frames) >= 1):
#nii to np
image_frames, head, affine = data.load_nii_image(
#with h5py.File('../dataset_impl/patches4/train.h5', 'r') as hf:
# data_train = np.array(hf.get('data'))
# label_train = np.array(hf.get('label'))
##split in and testset
data_train, label_train, data_test, label_test = minidataset.extract(
'../../dataset_impl/patches4', 20, 6, 15) #0,0
data_provider = image_util.SimpleDataProvider(data_train,
label_train,
channels_in=5,
channels_out=4,
n_class=16)
##setup & training
net = unet.Unet(channels_in=5, channels_out=4, n_class=16)
trainer = unet.Trainer(net, batch_size=1, optimizer="momentum") #10
path = trainer.train(data_provider, "prediction", training_iters=20,
epochs=6) #51-100
#verification
#prediction = net.predict(path, data_test) #data=testset
#unet.error_rate(prediction, util.crop_to_shape(label_test, prediction.shape))
#modified through reshape
#true_y=util.to_rgb(util.crop_to_shape(label_test, prediction.shape))
#est_y=util.to_rgb(prediction)
#util.save_image(true_y, 'true_y_fin.jpg')
val_generator = DataProviderMultitask(tf_records_file=tf_file_val,
sess=sess,
images2epoch=num_images_val *
aug_images_val,
IMAGE_H=nx,
IMAGE_W=ny,
batch_size=batch_size,
channels=1,
shuffle_flag=True,
augment=False)
net = unet.Unet(batch_size=batch_size,
IMAGE_H=nx,
IMAGE_W=ny,
channels=1,
n_class=num_classes,
cost="cross_entropy",
cost_kwargs=dict(class_weights=None),
root_filters=32,
branches=branches)
trainer_instance = seg_trainer.Trainer(
net, optimizer="adam", opt_kwargs=dict(learning_rate=0.001))
path = trainer_instance.train_val(
train_provider=train_generator,
val_provider=val_generator,
output_path=output_path,
sess=sess,
decay_step=10, ## decay every 10 epochs
epochs=100, # number of epochs to train
figure_path = path + "report/images/"
# load numpys
image_cn_ds = np.load(numpy_path + "image_vent_ds.npy")
image_nc_ds = np.load(numpy_path + "image_nc_ds.npy")
masks_ds = np.load(numpy_path + "masks_vent_ds.npy")
# images for explanation
image_cone = image_cn_ds[0, :, :, :]
image_no_cone = image_nc_ds[0, :, :, :]
mask = masks_ds[0, :, :, :]
# load model for lv seg
predict_image_params = params.Params(dim=[800, 800],
experiment_name="vent_cv_1")
model = unet.Unet(predict_image_params)
predicted_vents = model.make_prediction(image)
num_image_frames = image.shape[0]
# ROC
# image for tracking LVV though time
image = np.load(path + "data/single_img/image_nc.npy")
# LV vol through time
vent_size = np.empty([num_image_frames])
print(vent_size.shape)
for i in range(0, num_image_frames):
vent_size[i] = np.count_nonzero(predicted_vents[i, :, :, :])
start_time = time.time()
for combination in combinations[rank * perrank:(rank + 1) * perrank]:
stacks = combination[0]
stack_width = combination[1]
filters = combination[2]
kernel_size = (combination[3], combination[3])
checkpoint_path = "checkpoints/{}_{}_{}_{}/".format(
stacks, stack_width, filters, combination[3])
costs_path = "costs/{}_{}_{}_{}/".format(stacks, stack_width, filters,
combination[3])
model = unet.Unet(stacks,
stack_width,
filters,
output_channels,
kernel_size=kernel_size,
**unet_kw)
optimizer = tf.keras.optimizers.Adam(learning_rate)
for epoch in range(epochs):
batched_train_idx = rand.permutation(train_idx).reshape(
(-1, batch_size))
#Train
for iteration, snapshot_num in enumerate(batched_train_idx):
# Load adjascent outputs
inputs_0 = train_inputs[snapshot_num]
labels_0 = train_labels[snapshot_num]
patient_pre_list.append(image_pre_list)
patient_label_list.append(image_label_list)
patient_save_list.append(image_save_list)
print('start to predict')
if para.channel == 1:
generator = image_gen.oneChannelProvider(provider_path)
elif para.channel == 4:
generator = image_gen.fourChannelProvider(provider_path)
elif para.channel == 8:
generator = image_gen.eightChannelProvider(provider_path)
net = unet.Unet(channels=generator.channels,
n_class=generator.n_class,
cost=para.cost,
cost_kwargs=dict(regularizer=para.regularizer),
layers=para.layers,
features_root=para.features_root,
training=False)
def process_data(data):
# normalization
np.fabs(data)
data -= np.amin(data)
data /= np.amax(data)
return data
def process_labels(label):
nx = label.shape[1]
