def main():
# Parse command line args
if len(sys.argv) < 2:
print 'USAGE: binarize_image.py project_file model_name best_epoch_number test_image_file [models_dir]'
sys.exit(1)
project_file = sys.argv[1]
model_name = sys.argv[2]
best_epoch_number = int(sys.argv[3])
test_image_file = sys.argv[4]
if len(sys.argv) >= 6:
models_dir = sys.argv[5]
else:
models_dir = '.'
# Load the project file and the model
project = Project(project_file)
model_filename = os.path.join(models_dir, model_name + '.model')
model = get_model_from_name(project, model_name)
# Load the weight corresponding to the right epoch
load_weights(model, model_name, model_filename, best_epoch_number)
# Load the image
dataset = DataSet('.')
dataset.imagespaths[test_image_file] = test_image_file
data = TrainingData(dataset, dataset)
patch_size = project.get_training_dataset().config['patch_size']
data.set_config({'patch_size': patch_size})
genconfig = data.exhaustive_gen_config([patch_size, patch_size])
patches = data.generate_input_data(genconfig)
# Binarize the image
binarized_image = binarize(patches, patch_size, genconfig[test_image_file], model, dataset.open_image(test_image_file).size)
# Draw the image using PIL for easier image saving method
im = Image.fromarray(np.uint8(binarized_image*255))
im.show()
class Project(object):
def __init__(self, config_filepath):
self.config = json.load(open(config_filepath, 'r'))
self.training = None
self.training_genconfig = None
self.validation = None
self.validation_genconfig = None
self.models = [None] * len(self.config['models'])
self.load_gen_config = {
'random': lambda dataset, conf: dataset.random_gen_config(conf['patches_per_image']),
'exhaustive': lambda dataset, conf: dataset.exhaustive_gen_config(conf['patches_padding']),
'load': lambda dataset, conf: dataset.load_gen_config(['conf.generation_file'])
}
def get_training_dataset(self):
if self.training == None:
print 'Preloading training dataset...'
self.training = TrainingData(DataSet(self.config['training']['samples']), DataSet(self.config['training']['ground_truth']))
self.training.set_config({'patch_size': self.config['patch_size']})
return self.training
def get_training_genconfig(self):
if self.training_genconfig == None:
print 'Generating training dataset configuration...'
self.training_genconfig = self.load_gen_config.get(self.config['training']['generation_type'])(self.get_training_dataset(), self.config['training'])
return self.training_genconfig
def get_validation_dataset(self):
if self.validation == None:
print 'Generating validation dataset configuration...'
self.validation = TrainingData(DataSet(self.config['validation']['samples']), DataSet(self.config['validation']['ground_truth']))
self.validation.set_config({'patch_size': self.config['patch_size']})
return self.validation
def get_validation_genconfig(self):
if self.validation_genconfig == None:
print 'Generating validation dataset configuration...'
self.validation_genconfig = self.load_gen_config.get(self.config['validation']['generation_type'])(self.get_validation_dataset(), self.config['validation'])
return self.validation_genconfig
def get_models_count(self):
return len(self.models)
def get_model(self, i):
if self.models[i] == None:
model_config = self.config['models'][i]
print 'Configuring model ' + model_config['name'] + '...'
self.models[i] = NeuralNetwork()
patch_size = self.config['patch_size']
self.models[i].initialise(self.config['patch_size'] ** 2, range(0, model_config['network']['learning_params']['epochs'], model_config['logging_period']), model_config['network'])
return self.models[i]
def get_model_name(self, i):
return self.config['models'][i]['name']
def get_validation_dataset(self):
if self.validation == None:
print 'Generating validation dataset configuration...'
self.validation = TrainingData(DataSet(self.config['validation']['samples']), DataSet(self.config['validation']['ground_truth']))
self.validation.set_config({'patch_size': self.config['patch_size']})
return self.validation
def get_training_dataset(self):
if self.training == None:
print 'Preloading training dataset...'
self.training = TrainingData(DataSet(self.config['training']['samples']), DataSet(self.config['training']['ground_truth']))
self.training.set_config({'patch_size': self.config['patch_size']})
return self.training
