def test_read_object_classes(self):
map_filename = 'test_category_maps/stanford_bground_categories.txt'
category_colors, category_names = read_object_classes(map_filename)
self.assertEqual(category_colors[0], (int(0.3700 * 255), int(0.7300 * 255), int(0.9900 * 255)))
self.assertEqual(category_colors[1], (int(0.1500 * 255), int(0.3100 * 255), int(0.0700 * 255)))
self.assertEqual(category_colors[2], (int(0.2600 * 255), int(0.2600 * 255), int(0.2600 * 255)))
self.assertEqual(category_colors[3], (int(0.4200 * 255), int(0.6600 * 255), int(0.3100 * 255)))
self.assertEqual(category_colors[4], (int(0.0700 * 255), int(0.3300 * 255), int(0.8000 * 255)))
self.assertEqual(category_colors[5], (int(0.8800 * 255), int(0.4000 * 255), int(0.4000 * 255)))
self.assertEqual(category_colors[6], (int(0.8100 * 255), int(0.8900 * 255), int(0.9500 * 255)))
self.assertEqual(category_colors[7], (int(0.9500 * 255), int(0.7600 * 255), int(0.2000 * 255)))
category_names_correct = ["sky", "tree/bush", "road/path", "grass", "water", "building", "mountain",
"foreground_object"]
self.assertListEqual(category_names, category_names_correct)
def main():
# parse command line arguments
parser = argparse.ArgumentParser(
description='An rCNN scene labeling model.',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument(
'--dataset',
type=str,
default=FROM_GAMES,
choices=DATASETS.keys(),
help=
'Type of dataset to use. This determines the expected format of the data directory'
)
parser.add_argument('--data_dir',
type=str,
help='Directory for image and label data')
parser.add_argument('--category_map',
type=str,
help='File that maps colors ')
parser.add_argument('--hidden_size_1',
type=int,
default=25,
help='First Hidden size for CNN model')
parser.add_argument('--hidden_size_2',
type=int,
default=50,
help='Second Hidden size for CNN model')
parser.add_argument('--patch_size',
type=int,
default=67,
help='Patch size for input images')
parser.add_argument('--learning_rate',
type=float,
default=1e-4,
help='Learning rate for training CNN model')
# TODO figure out batch size
parser.add_argument('--batch_size',
type=int,
default=1,
help='Batch size for training CNN model')
parser.add_argument('--num_epochs',
type=int,
default=1,
help='Number of epochs for training CNN model')
parser.add_argument('--use_patches',
action='store_true',
default=False,
help='Whether to train model on individual patches')
parser.add_argument(
'--patches_per_image',
type=int,
default=1000,
help=
'Number of patches to sample for each image during training of CNN model'
)
parser.add_argument(
'--gaussian_sigma',
type=int,
choices=[15, 30],
default=None,
help='Size of gaussian mask to apply to patches. Not used by default.')
parser.add_argument(
'--fix_random_seed',
action='store_true',
default=False,
help='Whether to reset random seed at start, for debugging.')
parser.add_argument('--model_save_path',
type=str,
default=None,
help='Optional location to store saved model in.')
parser.add_argument('--model_load_path',
type=str,
default=None,
help='Optional location to load saved model from.')
parser.add_argument(
'--dry_run',
action='store_true',
default=False,
help=
'If true, only trains on one image, to test the training code quickly.'
)
parser.add_argument(
'--train_fraction',
type=float,
default=0.8,
help=
'Fraction of data to train on. If positive, trains on first X images, otherwise trains on '
'last X images.')
args = parser.parse_args()
if args.fix_random_seed:
random.seed(0)
# load class labels
category_colors, category_names, names_to_ids = read_object_classes(
args.category_map)
num_classes = len(category_names)
# create function that when called, provides iterator to an epoch of the data
dataset_func = DATASETS[args.dataset]
if args.dry_run:
def dataset_epoch_iter():
return dataset_func(args.data_dir, num_train=1)
else:
def dataset_epoch_iter():
return dataset_func(args.data_dir,
train_fraction=args.train_fraction)
model = CNNModel(args.hidden_size_1,
args.hidden_size_2,
args.batch_size,
num_classes,
args.learning_rate,
num_layers=2)
sess = tf.Session()
init = tf.initialize_all_variables()
sess.run(init)
if args.model_load_path is not None:
restore_model(sess, args.model_load_path)
train(sess,
model,
dataset_epoch_iter,
num_epochs=args.num_epochs,
use_patches=args.use_patches,
patches_per_image=args.patches_per_image,
save_path=args.model_save_path,
gaussian_sigma=args.gaussian_sigma)
print "Saving trained model to %s ..." % args.model_save_path
save_model(sess, args.model_save_path)
def main():
"""
Trains or evaluates an rCNN model.
"""
# parse command line arguments
parser = argparse.ArgumentParser(
description='An rCNN scene labeling model.',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument(
'--dataset',
type=str,
default=FROM_GAMES,
choices=DATASETS.keys(),
help=
'Type of dataset to use. This determines the expected format of the data directory'
)
parser.add_argument('--category_map',
type=str,
help='File that maps colors ')
parser.add_argument('--data_dir',
type=str,
help='Directory for image and label data')
parser.add_argument(
'--data_fraction',
type=float,
default=None,
help=
'Fraction of data to train on. If positive, trains on first X images, otherwise trains on '
'last X images.')
parser.add_argument('--hidden_size_1',
type=int,
default=10,
help='First Hidden size for CNN model')
parser.add_argument('--hidden_size_2',
type=int,
default=10,
help='Second Hidden size for CNN model')
parser.add_argument('--learning_rate',
type=float,
default=0.001,
help='Learning rate for training CNN model')
parser.add_argument('--num_epochs',
type=int,
default=1,
help='Number of epochs for training CNN model')
parser.add_argument('--model_save_path',
type=str,
default=None,
help='Optional location to store saved model in.')
parser.add_argument('--model_load_path',
type=str,
default=None,
help='Optional location to load saved model from.')
parser.add_argument('--training',
action='store_true',
default=False,
help='Whether or not to train model.')
parser.add_argument('--output_dir',
type=str,
default=None,
help='Directory in which to save test output images.')
parser.add_argument(
'--dry_run',
action='store_true',
default=False,
help=
'If true, only trains on one image, to test the training code quickly.'
)
parser.add_argument(
'--fix_random_seed',
action='store_true',
default=False,
help='Whether to reset random seed at start, for debugging.')
args = parser.parse_args()
if args.fix_random_seed:
random.seed(0)
# load class labels
category_colors, category_names = read_object_classes(args.category_map)
num_classes = len(category_names)
# create function that when called, provides iterator to an epoch of the data
dataset_func = DATASETS[args.dataset]
if args.dry_run:
def dataset_epoch_iter():
return dataset_func(args.data_dir, num_samples=1)
else:
def dataset_epoch_iter():
return dataset_func(args.data_dir,
data_fraction=args.data_fraction)
model = RCNNModel(args.hidden_size_1,
args.hidden_size_2,
num_classes,
args.learning_rate,
num_layers=2)
sess = tf.Session()
init = tf.initialize_all_variables()
sess.run(init)
if args.model_load_path is not None:
restore_model(sess, args.model_load_path)
run_model(sess,
model,
dataset_epoch_iter,
num_epochs=args.num_epochs,
training=args.training,
save_path=args.model_save_path,
output_dir=args.output_dir,
color_map=category_colors)
def main():
# parse command line arguments
parser = argparse.ArgumentParser(
description='Evaluate an rCNN scene labelling model.',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--model', type=str, help='Filename of saved model')
parser.add_argument('--category_map',
type=str,
help='File that maps colors ')
parser.add_argument(
'--dataset',
type=str,
default=FROM_GAMES,
choices=DATASETS.keys(),
help=
'Type of dataset to use. This determines the expected format of the data directory'
)
parser.add_argument('--data_dir',
type=str,
help='Directory for image and label data')
parser.add_argument(
'--output_dir',
type=str,
default=None,
help=
'Directory to store model output. By default no output is generated.')
parser.add_argument('--patch_size',
type=int,
default=67,
help='Size of input patches')
parser.add_argument('--use_patches',
action='store_true',
default=False,
help='Whether to evaluate model on individual patches')
parser.add_argument(
'--patches_per_image',
type=int,
default=2000,
help=
'Number of patches to sample from each test image. Not used by default.'
)
parser.add_argument(
'--gaussian_sigma',
type=int,
choices=[15, 30],
default=None,
help='Size of gaussian mask to apply to patches. Not used by default.')
parser.add_argument(
'--test_fraction',
type=float,
default=-0.2,
help=
'Fraction of data to test on. If positive, tests on first X images, otherwise tests on '
'last X images.')
parser.add_argument('--layer',
choices=[1, 2],
type=int,
default=2,
help='Number of rCNN layers to use.')
args = parser.parse_args()
# load class labels
category_colors, category_names, names_to_ids = read_object_classes(
args.category_map)
num_classes = len(category_names)
# load dataset
def dataset_func():
return DATASETS[args.dataset](args.data_dir,
train_fraction=args.test_fraction)
# TODO only test?
# TODO don't hardcode these (maybe store them in config file?)
model = CNNModel(25, 50, 1, num_classes, 1e-4, num_layers=2)
sess = tf.Session()
restore_model(sess, args.model)
test_model(sess,
model,
dataset_func,
args.layer,
use_patches=args.use_patches,
patches_per_image=args.patches_per_image,
gaussian_sigma=args.gaussian_sigma,
output_dir=args.output_dir,
color_map=category_colors)