def build_dwd_net(input,model,num_classes,pretrained_dir,substract_mean = False):
g, init_fn = build_refinenet(input, preset_model=model, num_classes=None, pretrained_dir=pretrained_dir,
substract_mean=substract_mean)
network_heads = dict()
with tf.variable_scope('deep_watershed'):
network_heads["stamp_class"] = dict()
# class binary
network_heads["stamp_class"]["binary"] = [slim.conv2d(g[x], 2, [1, 1], activation_fn=None, scope='class_binary_'+str(x)) for x in range(0, len(g))]
# class pred
network_heads["stamp_class"]["softmax"] = [slim.conv2d(g[x], num_classes, [1, 1], activation_fn=None, scope='class_pred_' + str(x)) for x in range(0, len(g))]
# direction
network_heads["stamp_directions"] = dict()
network_heads["stamp_directions"]["reg"] = [slim.conv2d(g[x], 2, [1, 1], activation_fn=None, scope='direction_' + str(x)) for x in range(0, len(g))]
network_heads["stamp_energy"] = dict()
# energy marker - regression
network_heads["stamp_energy"]["reg"] = [slim.conv2d(g[x], 1, [1, 1], activation_fn=None, scope='energy_reg_' + str(x)) for x in range(0, len(g))]
# energy marker - logits
network_heads["stamp_energy"]["softmax"] = [slim.conv2d(g[x], cfg.TRAIN.MAX_ENERGY, [1, 1], activation_fn=None, scope='energy_logits_' + str(x)) for x in range(0, len(g))]
network_heads["stamp_bbox"] = dict()
# bbox_size - reg
network_heads["stamp_bbox"]["reg"] = [slim.conv2d(g[x], 2, [1, 1], activation_fn=None, scope='bbox_reg_' + str(x)) for x in range(0, len(g))]
# bbox_size - logits
network_heads["stamp_bbox"]["softmax"] = [slim.conv2d(g[x], 2, [1, 1], activation_fn=None, scope='bbox_logits_' + str(x)) for x in range(0, len(g))]
return network_heads, init_fn
def inference (self, net_input, num_classes, is_training):
if FLAGS.patch_slim:
fuck_slim.patch(is_training)
network = None
init_fn = None
if FLAGS.net == "FC-DenseNet56" or FLAGS.net == "FC-DenseNet67" or FLAGS.net == "FC-DenseNet103":
with slim.arg_scope(aardvark.default_argscope(is_training)):
network = build_fc_densenet(net_input, preset_model = FLAGS.net, num_classes=num_classes)
elif FLAGS.net == "RefineNet-Res50" or FLAGS.net == "RefineNet-Res101" or FLAGS.net == "RefineNet-Res152":
with slim.arg_scope(aardvark.default_argscope(is_training)):
# RefineNet requires pre-trained ResNet weights
network, init_fn = build_refinenet(net_input, preset_model = FLAGS.net, num_classes=num_classes, is_training=is_training)
elif FLAGS.net == "FRRN-A" or FLAGS.net == "FRRN-B":
with slim.arg_scope(aardvark.default_argscope(is_training)):
network = build_frrn(net_input, preset_model = FLAGS.net, num_classes=num_classes)
elif FLAGS.net == "Encoder-Decoder" or FLAGS.net == "Encoder-Decoder-Skip":
with slim.arg_scope(aardvark.default_argscope(is_training)):
network = build_encoder_decoder(net_input, preset_model = FLAGS.net, num_classes=num_classes)
elif FLAGS.net == "MobileUNet" or FLAGS.net == "MobileUNet-Skip":
with slim.arg_scope(aardvark.default_argscope(is_training)):
network = build_mobile_unet(net_input, preset_model = FLAGS.net, num_classes=num_classes)
elif FLAGS.net == "PSPNet-Res50" or FLAGS.net == "PSPNet-Res101" or FLAGS.net == "PSPNet-Res152":
with slim.arg_scope(aardvark.default_argscope(is_training)):
# Image size is required for PSPNet
# PSPNet requires pre-trained ResNet weights
network, init_fn = build_pspnet(net_input, label_size=[args.crop_height, args.crop_width], preset_model = FLAGS.net, num_classes=num_classes, is_training=is_training)
elif FLAGS.net == "GCN-Res50" or FLAGS.net == "GCN-Res101" or FLAGS.net == "GCN-Res152":
with slim.arg_scope(aardvark.default_argscope(is_training)):
# GCN requires pre-trained ResNet weights
network, init_fn = build_gcn(net_input, preset_model = FLAGS.net, num_classes=num_classes, is_training=is_training)
elif FLAGS.net == "DeepLabV3-Res50" or FLAGS.net == "DeepLabV3-Res101" or FLAGS.net == "DeepLabV3-Res152":
with slim.arg_scope(aardvark.default_argscope(is_training)):
# DeepLabV requires pre-trained ResNet weights
network, init_fn = build_deeplabv3(net_input, preset_model = FLAGS.net, num_classes=num_classes, is_training=is_training)
elif FLAGS.net == "DeepLabV3_plus-Res50" or FLAGS.net == "DeepLabV3_plus-Res101" or FLAGS.net == "DeepLabV3_plus-Res152":
# DeepLabV3+ requires pre-trained ResNet weights
with slim.arg_scope(aardvark.default_argscope(is_training)):
network, init_fn = build_deeplabv3_plus(net_input, preset_model = FLAGS.net, num_classes=num_classes, is_training=is_training)
elif FLAGS.net == "AdapNet":
with slim.arg_scope(aardvark.default_argscope(is_training)):
network = build_adaptnet(net_input, num_classes=num_classes)
else:
raise ValueError("Error: the model %d is not available. Try checking which models are available using the command python main.py --help")
self.init_fn = init_fn
return network
# Compute your softmax cross entropy loss
print("Preparing the model ...")
net_input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
net_output = tf.placeholder(tf.float32, shape=[None, None, None, num_classes])
network = None
init_fn = None
if args.model == "FC-DenseNet56" or args.model == "FC-DenseNet67" or args.model == "FC-DenseNet103":
network = build_fc_densenet(net_input,
preset_model=args.model,
num_classes=num_classes)
elif args.model == "RefineNet-Res50" or args.model == "RefineNet-Res101" or args.model == "RefineNet-Res152":
# RefineNet requires pre-trained ResNet weights
network, init_fn = build_refinenet(net_input,
preset_model=args.model,
num_classes=num_classes)
elif args.model == "FRRN-A" or args.model == "FRRN-B":
network = build_frrn(net_input,
preset_model=args.model,
num_classes=num_classes)
elif args.model == "Encoder-Decoder" or args.model == "Encoder-Decoder-Skip":
network = build_encoder_decoder(net_input,
preset_model=args.model,
num_classes=num_classes)
elif args.model == "MobileUNet" or args.model == "MobileUNet-Skip":
network = build_mobile_unet(net_input,
preset_model=args.model,
num_classes=num_classes)
elif args.model == "PSPNet-Res50" or args.model == "PSPNet-Res101" or args.model == "PSPNet-Res152":
# Image size is required for PSPNet
print("Preparing the model ...")
input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
output = tf.placeholder(tf.float32, shape=[None, None, None, num_classes])
with tf.device('/gpu:' + str(args.gpu)):
input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
output = tf.placeholder(tf.float32, shape=[None, None, None, num_classes])
if args.model == "FC-DenseNet56" or args.model == "FC-DenseNet67" or args.model == "FC-DenseNet103" or args.model == "FC-DenseNet158" or args.model == "FC-DenseNet232":
network = build_fc_densenet(input,
preset_model=args.model,
num_classes=num_classes,
dropout_p=0.5)
elif args.model == "RefineNet-Res101" or args.model == "RefineNet-Res152":
network = build_refinenet(input,
preset_model=args.model,
num_classes=num_classes)
elif args.model == "Encoder-Decoder":
network = build_encoder_decoder(input, num_classes)
elif args.model == "Encoder-Decoder-Skip":
network = build_encoder_decoder_skip(input, num_classes)
elif args.model == "HF-FCN":
network = build_hf_fcn(input,
preset_model='FC-DenseNet103',
num_classes=num_classes)
elif args.model == "custom":
network = build_custom(input, num_classes)
# Compute your (unweighted) softmax cross entropy loss
if not args.balanced_weight:
loss = tf.reduce_mean(
def buildNetwork(model, net_input, num_class):
# Get the selected model.
# Some of them require pre-trained ResNet
if "Res50" in model and not os.path.isfile("models/resnet_v2_50.ckpt"):
download_checkpoints("ResnetV2", "50")
if "Res101" in model and not os.path.isfile("models/resnet_v2_101.ckpt"):
utils.download_checkpoints("ResnetV2", "101")
if "Res152" in model and not os.path.isfile("models/resnet_v2_152.ckpt"):
utils.download_checkpoints("ResnetV2", "152")
network = None
init_fn = None
if model == "FC-DenseNet56" or model == "FC-DenseNet67" or model == "FC-DenseNet103":
network = build_fc_densenet(net_input,
preset_model=model,
num_classes=num_class)
elif model == "RefineNet-Res50" or model == "RefineNet-Res101" or model == "RefineNet-Res152":
# RefineNet requires pre-trained ResNet weights
network, init_fn = build_refinenet(net_input,
preset_model=model,
num_classes=num_class)
elif model == "FRRN-A" or model == "FRRN-B":
network = build_frrn(net_input,
preset_model=model,
num_classes=num_class)
elif model == "Encoder-Decoder" or model == "Encoder-Decoder-Skip":
network = build_encoder_decoder(net_input,
preset_model=model,
num_classes=num_class)
elif model == "MobileUNet" or model == "MobileUNet-Skip":
network = build_mobile_unet(net_input,
preset_model=model,
num_classes=num_class)
elif model == "PSPNet-Res50" or model == "PSPNet-Res101" or model == "PSPNet-Res152":
# Image size is required for PSPNet
# PSPNet requires pre-trained ResNet weights
network, init_fn = build_pspnet(
net_input,
label_size=[args.crop_height, args.crop_width],
preset_model=model,
num_classes=num_class)
elif model == "GCN-Res50" or model == "GCN-Res101" or model == "GCN-Res152":
# GCN requires pre-trained ResNet weights
network, init_fn = build_gcn(net_input,
preset_model=model,
num_classes=num_class)
elif model == "DeepLabV3-Res50" or model == "DeepLabV3-Res101" or model == "DeepLabV3-Res152":
# DeepLabV requires pre-trained ResNet weights
network, init_fn = build_deeplabv3(net_input,
preset_model=model,
num_classes=num_class)
elif model == "DeepLabV3_plus-Res50" or model == "DeepLabV3_plus-Res101" or model == "DeepLabV3_plus-Res152":
# DeepLabV3+ requires pre-trained ResNet weights
network, init_fn = build_deeplabv3_plus(net_input,
preset_model=model,
num_classes=num_class)
elif model == "AdapNet":
network = build_adaptnet(net_input, num_classes=num_class)
elif model == "custom":
network = build_custom(net_input, num_class)
else:
raise ValueError(
"Error: the model %d is not available. Try checking which models are available using the command python main.py --help"
)
return network, init_fn
def main(unused_argv):
print("Setting up image reader...")
if args.dataset == "DeepScores":
data_reader = deepscores_semseg_datareader.ds_semseg_datareader(
cfg.DATA_DIR + "/DeepScores_2017/segmentation_detection",
crop=args.crop,
crop_size=args.crop_size)
#resnet_dir = cfg.PRETRAINED_DIR+"/DeepScores/"
resnet_dir = '../classification_pretrain/pretrain_deepscores/DeepScores/'
#refinenet_dir = cfg.PRETRAINED_DIR+"/DeepScores_semseg/"
refinenet_dir = os.getcwd()
num_classes = 124
input = tf.placeholder(
tf.float32, shape=[None, args.crop_size[0], args.crop_size[1], 1])
substract_mean = False
elif args.dataset == "VOC2012":
data_reader = pascalvoc_semseg_datareader.voc_seg_dataset_reader(
cfg.DATA_DIR + "/VOC2012")
resnet_dir = cfg.PRETRAINED_DIR + "/ImageNet/"
#refinenet_dir = cfg.PRETRAINED_DIR+"/VOC2012/"
refinenet_dir = os.getcwd()
num_classes = 21
input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
substract_mean = True
else:
print("Unknown dataset")
sys.exit(1)
#train_images, train_annotations = data_reader.next_batch(20)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
# init model
print("Preparing the model ...")
output = tf.placeholder(tf.float32, shape=[None, None, None, num_classes])
network, init_fn = build_refinenet(input,
preset_model=args.model,
num_classes=num_classes,
pretrained_dir=resnet_dir,
substract_mean=substract_mean)
# init optimizer
# Compute your (unweighted) softmax cross entropy loss
loss = tf.reduce_mean(
tf.nn.softmax_cross_entropy_with_logits(logits=network, labels=output))
opt = tf.train.RMSPropOptimizer(
learning_rate=0.0001, decay=0.995).minimize(
loss, var_list=[var for var in tf.trainable_variables()])
saver = tf.train.Saver(max_to_keep=1000)
sess.run(tf.global_variables_initializer())
# load classification weights or continue training from older checkpoint
# If a pre-trained ResNet is required, load the weights. --> depends on path tiven to build_refinenet
# This must be done AFTER the variables are initialized with sess.run(tf.global_variables_initializer())
print("loading resnet-weights")
if init_fn is not None:
init_fn(sess)
# Load a previous checkpoint if desired
model_checkpoint_name = refinenet_dir + args.model + ".ckpt"
if args.continue_training:
print('Loaded latest model checkpoint')
saver.restore(sess, model_checkpoint_name)
print("start training")
# Train for a fixed number of batches
for itr in range(1, (args.iterations + 1)):
train_images, train_annotations = data_reader.next_batch(
args.batch_size)
# VOC only works with batchsize 1
if args.dataset == "VOC2012":
train_images = np.expand_dims(train_images[0], 0)
train_annotations = np.expand_dims(train_annotations[0], 0)
# for i in range(0,train_images.shape[0]):
# train_annotations[i] = np.eye(num_classes)[train_annotations[i][:,:,-1]]
# convert annotation to one-hot enc
train_annotations = np.eye(num_classes)[train_annotations[:, :, :, -1]]
# size test
# h = 400
# w = 500
# train_images = np.zeros((1,h,w,3))
# train_annotations = np.zeros((1,h,w,21))
# h and w has to be a multiple of 32!!
# gets taken care of by data loaders
_, current = sess.run([opt, loss],
feed_dict={
input: train_images,
output: train_annotations
})
if itr == 1:
print("initial loss" + str(current))
if itr % 21 == 0:
print("loss of current batch:" + str(current))
if itr % 2001 == 0:
print("saving weights")
saver.save(sess, model_checkpoint_name)
print("Compute test error")
test_images, test_annotations = data_reader.get_test_records()
test_annotations = np.eye(num_classes)[test_annotations[:, :, :, -1]]
test_losses = []
for itr in range(0, len(test_images) / args.batch_size):
test_img_b = test_images[(itr * args.batch_size):((itr + 1) *
args.batch_size)]
test_ann_b = test_annotations[(itr *
args.batch_size):((itr + 1) *
args.batch_size)]
cur_loss = sess.run([loss],
feed_dict={
input: test_img_b,
output: test_ann_b
})
test_losses.append(cur_loss)
print("Test Loss:")
print(np.mean(test_losses))
def build_model(model_name,
net_input,
num_classes,
frontend="ResNet101",
is_training=True):
# Get the selected model.
# Some of them require pre-trained ResNet
print("Preparing the model ...")
if model_name not in SUPPORTED_MODELS:
raise ValueError(
"The model you selelect is not supported. The following models are currently supported: {0}"
.format(SUPPORTED_MODELS))
if frontend not in SUPPORTED_FRONTENDS:
raise ValueError(
"The frontend you selelect is not supported. The following models are currently supported: {0}"
.format(SUPPORTED_FRONTENDS))
if "ResNet50" == frontend and not os.path.isfile(
"models/resnet_v2_50.ckpt"):
download_checkpoints("ResNet50")
if "ResNet101" == frontend and not os.path.isfile(
"models/resnet_v2_101.ckpt"):
download_checkpoints("ResNet101")
if "ResNet152" == frontend and not os.path.isfile(
"models/resnet_v2_152.ckpt"):
download_checkpoints("ResNet152")
if "MobileNetV2" == frontend and not os.path.isfile(
"models/mobilenet_v2_1.4_224.ckpt.data-00000-of-00001"):
download_checkpoints("MobileNetV2")
if "InceptionV4" == frontend and not os.path.isfile(
"models/inception_v4.ckpt"):
download_checkpoints("InceptionV4")
network = None
init_fn = None
if model_name == "FC-DenseNet56" or model_name == "FC-DenseNet67" or model_name == "FC-DenseNet103":
network = build_fc_densenet(net_input,
preset_model=model_name,
num_classes=num_classes)
elif model_name == "RefineNet":
# RefineNet requires pre-trained ResNet weights
network, init_fn = build_refinenet(net_input,
preset_model=model_name,
frontend=frontend,
num_classes=num_classes,
is_training=is_training)
elif model_name == "FRRN-A" or model_name == "FRRN-B":
network = build_frrn(net_input,
preset_model=model_name,
num_classes=num_classes)
elif model_name == "Encoder-Decoder" or model_name == "Encoder-Decoder-Skip":
network = build_encoder_decoder(net_input,
preset_model=model_name,
num_classes=num_classes)
elif model_name == "MobileUNet" or model_name == "MobileUNet-Skip":
network = build_mobile_unet(net_input,
preset_model=model_name,
num_classes=num_classes)
elif model_name == "PSPNet":
# Image size is required for PSPNet
# PSPNet requires pre-trained ResNet weights
network, init_fn = build_pspnet(
net_input,
label_size=[args.crop_height, args.crop_width],
preset_model=model_name,
frontend=frontend,
num_classes=num_classes,
is_training=is_training)
elif model_name == "GCN":
# GCN requires pre-trained ResNet weights
network, init_fn = build_gcn(net_input,
preset_model=model_name,
frontend=frontend,
num_classes=num_classes,
is_training=is_training)
elif model_name == "DeepLabV3":
# DeepLabV requires pre-trained ResNet weights
network, init_fn = build_deeplabv3(net_input,
preset_model=model_name,
frontend=frontend,
num_classes=num_classes,
is_training=is_training)
elif model_name == "DeepLabV3_plus":
# DeepLabV3+ requires pre-trained ResNet weights
network, init_fn = build_deeplabv3_plus(net_input,
preset_model=model_name,
frontend=frontend,
num_classes=num_classes,
is_training=is_training)
elif model_name == "AdapNet":
network = build_adaptnet(net_input, num_classes=num_classes)
elif model_name == "custom":
network = build_custom(net_input, num_classes)
else:
raise ValueError(
"Error: the model %d is not available. Try checking which models are available using the command python main.py --help"
)
return network, init_fn
