def infer(self, input_image_path, return_image_paths=False, num_output=None):
"""Given an image, a path containing images, or a list of paths, return the outputs."""
one_output = False
if input_image_path:
if isinstance(input_image_path, list) or isinstance(input_image_path, tuple):
image_paths = input_image_path
else:
if os.path.isfile(input_image_path):
image_paths = [input_image_path]
one_output = True
else:
image_paths = util_io.get_files_in_dir(
input_image_path, do_sort=True, do_random_ordering=False)
images = [util_io.imread(image_path, dtype=np.uint8)
for image_path in image_paths]
else:
assert num_output >= 1
images = [None for _ in range(num_output)]
image_paths = [str(i) for i in range(num_output)]
one_output = (num_output == 1)
outputs = []
for image in images:
if image is None:
feed_dict = None
else:
feed_dict = {self.images_placeholder: image}
output = self.sess.run(self.output, feed_dict=feed_dict)
output = output[0] * 255.0 # Batch size == 1, range = 0~1.
outputs.append(output)
if one_output:
outputs = outputs[0]
image_paths = image_paths[0]
if return_image_paths:
return outputs, image_paths
return outputs
def _find_image_files(self, data_dir, shared_info):
file_names = util_io.get_files_in_dir(data_dir,
do_sort=False,
do_random_ordering=True)
if FLAGS.exclude_affix:
file_names = [
file_name for file_name in file_names
if not file_name.endswith(FLAGS.exclude_affix)
]
return file_names, [None for _ in range(len(file_names))]
def test_get_files_in_dir_in_dir(self):
dirpath = tempfile.mkdtemp()
image_path = dirpath + '/image.jpg'
f = open(image_path, 'w')
f.close()
subfolder = '/subfolder'
os.makedirs(dirpath + subfolder)
image_path2 = dirpath + subfolder + u'/骨董屋・三千世界の女主人_12746957.png'
f = open(image_path2, 'w')
f.close()
actual_answer = util_io.get_files_in_dir(dirpath + '/')
expected_answer = [image_path, image_path2.encode('utf8')]
shutil.rmtree(dirpath)
self.assertEqual(expected_answer, actual_answer)
def test_imread_and_imsave_utf8(self):
height = 256
width = 256
content_folder = tempfile.mkdtemp()
image_path = content_folder + u'/骨董屋・三千世界の女主人_12746957.png'
current_image = np.ones((height, width, 3)) * 255.0
current_image[0, 0, 0] = 0
util_io.imsave(image_path, current_image)
actual_output = util_io.imread(
util_io.get_files_in_dir(content_folder + '/')[0])
expected_answer = np.round(np.array(current_image))
np.testing.assert_almost_equal(expected_answer, actual_output)
def _generate_shared_info(self, train_directory, validation_directory,
output_path):
# Read TfExamples.
tf_records = util_io.get_files_in_dir(train_directory,
do_sort=True,
do_random_ordering=False,
allowed_extensions=None)
model = AnimeFaceObjectExtractor(tf_records,
do_safe_only=FLAGS.do_safe_only,
do_unsafe_only=FLAGS.do_unsafe_only)
model.main()
base_file_names = model.ret.keys()
ret = {
'object_detection_result_dict': model.ret,
'base_file_names': base_file_names
}
return ret
