def generate_and_save(self, data_dir, num_images):
"""
Generate and saves image in data directory
"""
data_dir = utils.make_valid_dir(data_dir)
try:
os.mkdir(data_dir)
except FileExistsError:
pass
# color_dict = {0:'Red', 1:'Green', 2:'Blue'}
color_dict = {0: 'Blue', 1: 'Green', 2: 'Red'} # For cv2
gen = self.generator()
for key in color_dict:
try:
os.mkdir(data_dir + color_dict[key])
except FileExistsError:
pass
for _ in tqdm(range(num_images // self.batch_size)):
batch = next(gen)
for count in range(self.batch_size):
img = batch[0][count]
label = color_dict[np.argmax(batch[1][count])]
cv2.imwrite(
data_dir + label + '/' + utils.random_string() + '.jpg',
img)
def transfer(self, model_dir):
"""
Tune our own classifier network over
pretrained feature extractor
"""
model_dir = utils.make_valid_dir(model_dir)
self.mode = 'transfer'
self.train_estimator(model_dir)
def input_fn():
"""
Input function
"""
data_dir = utils.make_valid_dir(self.data_dir)
list_of_data_path = [
data_dir + tfrecord for tfrecord in os.listdir(data_dir)
]
features, labels = self.tfrecord_iterator(list_of_data_path)
return features, labels
def retrain(self, model_dir):
"""
Retrain model graph
"""
model_dir = utils.make_valid_dir(model_dir)
self.mode = 'retrain'
assert (tf.gfile.Exists(model_dir)),\
"Model directory: '" + str(model_dir) + "' doesn't exists,\n" +\
"Train model from scratch before retraining"
self.train_estimator(model_dir)
def __init__(self,
data_dir,
bucket_name,
num_workers,
image_size=None,
num_classes=None,
batch_size=64, num_epochs=1,
mode='train'):
"""
Initialization
"""
self.data_dir = utils.make_valid_dir(data_dir)
self.bucket_name = bucket_name
self.image_size = image_size
self.num_classes = num_classes
self.batch_size = batch_size
self.num_epochs = num_epochs
self.mode = mode
self.list_of_data_path = self.make_list_of_data_path(
self.data_dir + utils.make_valid_dir(self.mode))
self.num_workers = num_workers
def convert(self, image_paths, labels, out_dir):
"""
Convert images to tfRecords and write them to out_dir
"""
try:
os.mkdir(out_dir)
except FileExistsError:
print('Warning: The output directory already exists')
out_dir = utils.make_valid_dir(out_dir)
total_images = len(image_paths)
index = 0
while index + self.batch_size <= total_images:
image_paths_batch = image_paths[index:index + self.batch_size]
labels_batch = labels[index:index + self.batch_size]
out_path = out_dir + 'batch-' + \
str(index / self.batch_size).zfill(12) + '.tfrecords'
self.convert_batch(image_paths_batch, labels_batch, out_path)
print('\rSerialized batch: ' + str(index // self.batch_size),
end='')
index += self.batch_size