diff --git a/.gitignore b/.gitignore index ab31201c..b6361af9 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,6 @@ *.pyc +*.graph +models/*.graph models/*.pyc models/*.ckpt* models/*.pb diff --git a/builders/__init__.py b/builders/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/models/frontend_builder.py b/builders/frontend_builder.py similarity index 96% rename from models/frontend_builder.py rename to builders/frontend_builder.py index 9ac97a47..c7a70725 100644 --- a/models/frontend_builder.py +++ b/builders/frontend_builder.py @@ -1,8 +1,8 @@ import tensorflow as tf from tensorflow.contrib import slim -import resnet_v2 -import mobilenet_v2 -import inception_v4 +from frontends import resnet_v2 +from frontends import mobilenet_v2 +from frontends import inception_v4 import os diff --git a/model_builder.py b/builders/model_builder.py similarity index 80% rename from model_builder.py rename to builders/model_builder.py index 31761c46..07cd5fc8 100644 --- a/model_builder.py +++ b/builders/model_builder.py @@ -3,16 +3,17 @@ import subprocess sys.path.append("models") -from FC_DenseNet_Tiramisu import build_fc_densenet -from Encoder_Decoder import build_encoder_decoder -from RefineNet import build_refinenet -from FRRN import build_frrn -from MobileUNet import build_mobile_unet -from PSPNet import build_pspnet -from GCN import build_gcn -from DeepLabV3 import build_deeplabv3 -from DeepLabV3_plus import build_deeplabv3_plus -from AdapNet import build_adaptnet +from models.FC_DenseNet_Tiramisu import build_fc_densenet +from models.Encoder_Decoder import build_encoder_decoder +from models.RefineNet import build_refinenet +from models.FRRN import build_frrn +from models.MobileUNet import build_mobile_unet +from models.PSPNet import build_pspnet +from models.GCN import build_gcn +from models.DeepLabV3 import build_deeplabv3 +from models.DeepLabV3_plus import build_deeplabv3_plus +from models.AdapNet import build_adaptnet +from models.custom_model import build_custom SUPPORTED_MODELS = ["FC-DenseNet56", "FC-DenseNet67", "FC-DenseNet103", "Encoder-Decoder", "Encoder-Decoder-Skip", "RefineNet", "FRRN-A", "FRRN-B", "MobileUNet", "MobileUNet-Skip", "PSPNet", "GCN", "DeepLabV3", "DeepLabV3_plus", "AdapNet", "custom"] @@ -20,11 +21,11 @@ SUPPORTED_FRONTENDS = ["ResNet50", "ResNet101", "ResNet152", "MobileNetV2", "InceptionV4"] def download_checkpoints(model_name): - subprocess.check_output(["python", "get_pretrained_checkpoints.py", "--model=" + model_name]) + subprocess.check_output(["python", "utils/get_pretrained_checkpoints.py", "--model=" + model_name]) -def build_model(model_name, net_input, num_classes, frontend="ResNet101", is_training=True): +def build_model(model_name, net_input, num_classes, crop_width, crop_height, frontend="ResNet101", is_training=True): # Get the selected model. # Some of them require pre-trained ResNet @@ -63,7 +64,7 @@ def build_model(model_name, net_input, num_classes, frontend="ResNet101", is_tra 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) + network, init_fn = build_pspnet(net_input, label_size=[crop_height, 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) diff --git a/frontends/__init__.py b/frontends/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/models/conv_blocks.py b/frontends/conv_blocks.py similarity index 100% rename from models/conv_blocks.py rename to frontends/conv_blocks.py diff --git a/models/inception_utils.py b/frontends/inception_utils.py similarity index 100% rename from models/inception_utils.py rename to frontends/inception_utils.py diff --git a/models/inception_v4.py b/frontends/inception_v4.py similarity index 99% rename from models/inception_v4.py rename to frontends/inception_v4.py index 2f9e5c01..e635fd6c 100644 --- a/models/inception_v4.py +++ b/frontends/inception_v4.py @@ -26,7 +26,7 @@ import tensorflow as tf -import inception_utils +from frontends import inception_utils slim = tf.contrib.slim diff --git a/models/mobilenet_base.py b/frontends/mobilenet_base.py similarity index 100% rename from models/mobilenet_base.py rename to frontends/mobilenet_base.py diff --git a/models/mobilenet_v2.py b/frontends/mobilenet_v2.py similarity index 98% rename from models/mobilenet_v2.py rename to frontends/mobilenet_v2.py index 23dac273..bc5e5060 100644 --- a/models/mobilenet_v2.py +++ b/frontends/mobilenet_v2.py @@ -29,8 +29,8 @@ import tensorflow as tf -import conv_blocks as ops -import mobilenet_base as lib +from frontends import conv_blocks as ops +from frontends import mobilenet_base as lib slim = tf.contrib.slim op = lib.op diff --git a/models/resnet_utils.py b/frontends/resnet_utils.py similarity index 100% rename from models/resnet_utils.py rename to frontends/resnet_utils.py diff --git a/models/resnet_v1.py b/frontends/resnet_v1.py similarity index 99% rename from models/resnet_v1.py rename to frontends/resnet_v1.py index 220bbd30..89a1d0d8 100644 --- a/models/resnet_v1.py +++ b/frontends/resnet_v1.py @@ -1,6 +1,6 @@ import tensorflow as tf from tensorflow.contrib import slim -import resnet_utils +from frontends import resnet_utils resnet_arg_scope = resnet_utils.resnet_arg_scope diff --git a/models/resnet_v2.py b/frontends/resnet_v2.py similarity index 99% rename from models/resnet_v2.py rename to frontends/resnet_v2.py index f6d86839..ee7550c8 100644 --- a/models/resnet_v2.py +++ b/frontends/resnet_v2.py @@ -52,7 +52,7 @@ import tensorflow as tf -import resnet_utils +from frontends import resnet_utils slim = tf.contrib.slim resnet_arg_scope = resnet_utils.resnet_arg_scope diff --git a/iou_vs_epochs.png b/iou_vs_epochs.png index 7937761e..3149dff0 100644 Binary files a/iou_vs_epochs.png and b/iou_vs_epochs.png differ diff --git a/models/AdapNet.py b/models/AdapNet.py index a5bcaf03..0fc03baa 100644 --- a/models/AdapNet.py +++ b/models/AdapNet.py @@ -3,7 +3,7 @@ import tensorflow as tf from tensorflow.contrib import slim import numpy as np -import resnet_v2 +from frontends import resnet_v2 import os, sys diff --git a/models/DeepLabV3.py b/models/DeepLabV3.py index 008a3f69..e7ecc115 100644 --- a/models/DeepLabV3.py +++ b/models/DeepLabV3.py @@ -3,7 +3,7 @@ import tensorflow as tf from tensorflow.contrib import slim import numpy as np -import frontend_builder +from builders import frontend_builder import os, sys def Upsampling(inputs,feature_map_shape): diff --git a/models/DeepLabV3_plus.py b/models/DeepLabV3_plus.py index 66dac5e1..f0489a8c 100644 --- a/models/DeepLabV3_plus.py +++ b/models/DeepLabV3_plus.py @@ -2,9 +2,8 @@ import tensorflow as tf from tensorflow.contrib import slim -import frontend_builder +from builders import frontend_builder import numpy as np -import resnet_v2 import os, sys def Upsampling(inputs,feature_map_shape): diff --git a/models/FRRN.py b/models/FRRN.py index 633186f2..80c38a63 100644 --- a/models/FRRN.py +++ b/models/FRRN.py @@ -1,6 +1,5 @@ import tensorflow as tf from tensorflow.contrib import slim -import resnet_v1 def Upsampling(inputs,scale): return tf.image.resize_nearest_neighbor(inputs, size=[tf.shape(inputs)[1]*scale, tf.shape(inputs)[2]*scale]) diff --git a/models/GCN.py b/models/GCN.py index 6c1c76e6..b2d98b51 100644 --- a/models/GCN.py +++ b/models/GCN.py @@ -1,6 +1,6 @@ import tensorflow as tf from tensorflow.contrib import slim -import frontend_builder +from builders import frontend_builder import os, sys def Upsampling(inputs,scale): diff --git a/models/ICNet.py b/models/ICNet.py index a6389e91..eb20d3ba 100644 --- a/models/ICNet.py +++ b/models/ICNet.py @@ -1,7 +1,7 @@ import tensorflow as tf from tensorflow.contrib import slim import numpy as np -import frontend_builder +from frontends import frontend_builder import os, sys def Upsampling_by_shape(inputs, feature_map_shape): diff --git a/models/PSPNet.py b/models/PSPNet.py index 1650a3f6..eb5ee81a 100644 --- a/models/PSPNet.py +++ b/models/PSPNet.py @@ -1,7 +1,7 @@ import tensorflow as tf from tensorflow.contrib import slim import numpy as np -import frontend_builder +from builders import frontend_builder import os, sys def Upsampling(inputs,feature_map_shape): diff --git a/models/RefineNet.py b/models/RefineNet.py index a9eae593..8636bb3a 100644 --- a/models/RefineNet.py +++ b/models/RefineNet.py @@ -1,6 +1,6 @@ import tensorflow as tf from tensorflow.contrib import slim -import frontend_builder +from builders import frontend_builder import os, sys def Upsampling(inputs,scale): @@ -142,7 +142,7 @@ def RefineBlock(high_inputs=None,low_inputs=None): -def build_refinenet(inputs, num_classes, preset_model='RefineNet', frontend="ResNet101", weight_decay=1e-5, upscaling_method="bilinear", pretrained_dir="models"): +def build_refinenet(inputs, num_classes, preset_model='RefineNet', frontend="ResNet101", weight_decay=1e-5, upscaling_method="bilinear", pretrained_dir="models", is_training=True): """ Builds the RefineNet model. diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/models/custom_model.py b/models/custom_model.py index 0362a582..83548611 100644 --- a/models/custom_model.py +++ b/models/custom_model.py @@ -3,7 +3,7 @@ import tensorflow as tf import tensorflow.contrib.slim as slim import numpy as np -import frontend_builder +from builders import frontend_builder def conv_block(inputs, n_filters, filter_size=[3, 3], dropout_p=0.0): """ diff --git a/predict.py b/predict.py index 69e38e58..aceac667 100644 --- a/predict.py +++ b/predict.py @@ -3,9 +3,8 @@ import argparse import numpy as np -import utils -import helpers -import model_builder +from utils import utils, helpers +from builders import model_builder parser = argparse.ArgumentParser() parser.add_argument('--image', type=str, default=None, required=True, help='The image you want to predict on. ') diff --git a/test.py b/test.py index 7cf8f42b..2d1de393 100644 --- a/test.py +++ b/test.py @@ -3,9 +3,8 @@ import argparse import numpy as np -import utils -import helpers -import model_builder +from utils import utils, helpers +from builders import model_builder parser = argparse.ArgumentParser() parser.add_argument('--checkpoint_path', type=str, default=None, required=True, help='The path to the latest checkpoint weights for your model.') diff --git a/train.py b/train.py index ddbf5b8d..29e6f4b4 100644 --- a/train.py +++ b/train.py @@ -10,14 +10,11 @@ import subprocess -import helpers -import utils +from utils import utils, helpers +from builders import model_builder import matplotlib.pyplot as plt -sys.path.append("models") -import model_builder - def str2bool(v): if v.lower() in ('yes', 'true', 't', 'y', '1'): return True @@ -90,7 +87,7 @@ def data_augmentation(input_image, output_image): net_input = tf.placeholder(tf.float32,shape=[None,None,None,3]) net_output = tf.placeholder(tf.float32,shape=[None,None,None,num_classes]) -network, init_fn = model_builder.build_model(args.model, frontend=args.frontend, net_input=net_input, num_classes=num_classes, is_training=True) +network, init_fn = model_builder.build_model(model_name=args.model, frontend=args.frontend, net_input=net_input, num_classes=num_classes, crop_width=args.crop_width, crop_height=args.crop_height, is_training=True) loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=network, labels=net_output)) diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/get_pretrained_checkpoints.py b/utils/get_pretrained_checkpoints.py similarity index 100% rename from get_pretrained_checkpoints.py rename to utils/get_pretrained_checkpoints.py diff --git a/helpers.py b/utils/helpers.py similarity index 100% rename from helpers.py rename to utils/helpers.py diff --git a/utils.py b/utils/utils.py similarity index 99% rename from utils.py rename to utils/utils.py index 8655380a..c476de00 100644 --- a/utils.py +++ b/utils/utils.py @@ -11,7 +11,7 @@ recall_score, confusion_matrix, classification_report, \ accuracy_score, f1_score -import helpers +from utils import helpers def prepare_data(dataset_dir): train_input_names=[]