def testWriteAndReadToFileEmptyFile(self):
filename = os.path.join(FLAGS.test_tmpdir, 'test.box')
box_io.WriteToFile(filename, np.array([]), np.array([]), np.array([]))
data_read = box_io.ReadFromFile(filename)
self.assertAllEqual(np.array([]), data_read[0])
self.assertAllEqual(np.array([]), data_read[1])
self.assertAllEqual(np.array([]), data_read[2])
def testWriteAndReadToFileEmptyFile(self):
tmpdir = tf.compat.v1.test.get_temp_dir()
filename = os.path.join(tmpdir, 'test.box')
box_io.WriteToFile(filename, np.array([]), np.array([]), np.array([]))
data_read = box_io.ReadFromFile(filename)
self.assertAllEqual(np.array([]), data_read[0])
self.assertAllEqual(np.array([]), data_read[1])
self.assertAllEqual(np.array([]), data_read[2])
def testWriteAndReadToFile(self):
boxes, scores, class_indices = self._create_data()
filename = os.path.join(FLAGS.test_tmpdir, 'test.boxes')
box_io.WriteToFile(filename, boxes, scores, class_indices)
data_read = box_io.ReadFromFile(filename)
self.assertAllEqual(boxes, data_read[0])
self.assertAllEqual(scores, data_read[1])
self.assertAllEqual(class_indices, data_read[2])
def ExtractBoxesAndFeaturesToFiles(image_names, image_paths, delf_config_path,
detector_model_dir, detector_thresh,
output_features_dir, output_boxes_dir,
output_mapping):
"""Extracts boxes and features, saving them to files.
Boxes are saved to <image_name>.boxes files. DELF features are extracted for
the entire image and saved into <image_name>.delf files. In addition, DELF
features are extracted for each high-confidence bounding box in the image, and
saved into files named <image_name>_0.delf, <image_name>_1.delf, etc.
It checks if descriptors/boxes already exist, and skips computation for those.
Args:
image_names: List of image names. These are used to compose output file
names for boxes and features.
image_paths: List of image paths. image_paths[i] is the path for the image
named by image_names[i]. `image_names` and `image_paths` must have the
same number of elements.
delf_config_path: Path to DelfConfig proto text file.
detector_model_dir: Directory where detector SavedModel is located.
detector_thresh: Threshold used to decide if an image's detected box
undergoes feature extraction.
output_features_dir: Directory where DELF features will be written to.
output_boxes_dir: Directory where detected boxes will be written to.
output_mapping: CSV file which maps each .delf file name to the image ID and
detected box ID.
Raises:
ValueError: If len(image_names) and len(image_paths) are different.
"""
num_images = len(image_names)
if len(image_paths) != num_images:
raise ValueError(
'image_names and image_paths have different number of items')
# Parse DelfConfig proto.
config = delf_config_pb2.DelfConfig()
with tf.io.gfile.GFile(delf_config_path, 'r') as f:
text_format.Merge(f.read(), config)
# Create output directories if necessary.
if not tf.io.gfile.exists(output_features_dir):
tf.io.gfile.makedirs(output_features_dir)
if not tf.io.gfile.exists(output_boxes_dir):
tf.io.gfile.makedirs(output_boxes_dir)
if not tf.io.gfile.exists(os.path.dirname(output_mapping)):
tf.io.gfile.makedirs(os.path.dirname(output_mapping))
names_ids_and_boxes = []
with tf.Graph().as_default():
with tf.compat.v1.Session() as sess:
# Initialize variables, construct detector and DELF extractor.
init_op = tf.compat.v1.global_variables_initializer()
sess.run(init_op)
detector_fn = detector.MakeDetector(sess,
detector_model_dir,
import_scope='detector')
delf_extractor_fn = extractor.MakeExtractor(
sess, config, import_scope='extractor_delf')
start = time.clock()
for i in range(num_images):
if i == 0:
print('Starting to extract features/boxes...')
elif i % _STATUS_CHECK_ITERATIONS == 0:
elapsed = (time.clock() - start)
print('Processing image %d out of %d, last %d '
'images took %f seconds' %
(i, num_images, _STATUS_CHECK_ITERATIONS, elapsed))
start = time.clock()
image_name = image_names[i]
output_feature_filename_whole_image = os.path.join(
output_features_dir, image_name + _DELF_EXTENSION)
output_box_filename = os.path.join(output_boxes_dir,
image_name + _BOX_EXTENSION)
pil_im = utils.RgbLoader(image_paths[i])
width, height = pil_im.size
# Extract and save boxes.
if tf.io.gfile.exists(output_box_filename):
print('Skipping box computation for %s' % image_name)
(boxes_out, scores_out, class_indices_out
) = box_io.ReadFromFile(output_box_filename)
else:
(boxes_out, scores_out, class_indices_out) = detector_fn(
np.expand_dims(pil_im, 0))
# Using only one image per batch.
boxes_out = boxes_out[0]
scores_out = scores_out[0]
class_indices_out = class_indices_out[0]
box_io.WriteToFile(output_box_filename, boxes_out,
scores_out, class_indices_out)
# Select boxes with scores greater than threshold. Those will be the
# ones with extracted DELF features (besides the whole image, whose DELF
# features are extracted in all cases).
num_delf_files = 1
selected_boxes = []
for box_ind, box in enumerate(boxes_out):
if scores_out[box_ind] >= detector_thresh:
selected_boxes.append(box)
num_delf_files += len(selected_boxes)
# Extract and save DELF features.
for delf_file_ind in range(num_delf_files):
if delf_file_ind == 0:
box_name = image_name
output_feature_filename = output_feature_filename_whole_image
else:
box_name = image_name + '_' + str(delf_file_ind - 1)
output_feature_filename = os.path.join(
output_features_dir, box_name + _DELF_EXTENSION)
names_ids_and_boxes.append(
[box_name, i, delf_file_ind - 1])
if tf.io.gfile.exists(output_feature_filename):
print('Skipping DELF computation for %s' % box_name)
continue
if delf_file_ind >= 1:
bbox_for_cropping = selected_boxes[delf_file_ind - 1]
bbox_for_cropping_pil_convention = [
int(math.floor(bbox_for_cropping[1] * width)),
int(math.floor(bbox_for_cropping[0] * height)),
int(math.ceil(bbox_for_cropping[3] * width)),
int(math.ceil(bbox_for_cropping[2] * height))
]
pil_cropped_im = pil_im.crop(
bbox_for_cropping_pil_convention)
im = np.array(pil_cropped_im)
else:
im = np.array(pil_im)
extracted_features = delf_extractor_fn(im)
locations_out = extracted_features['local_features'][
'locations']
descriptors_out = extracted_features['local_features'][
'descriptors']
feature_scales_out = extracted_features['local_features'][
'scales']
attention_out = extracted_features['local_features'][
'attention']
feature_io.WriteToFile(output_feature_filename,
locations_out, feature_scales_out,
descriptors_out, attention_out)
# Save mapping from output DELF name to image id and box id.
_WriteMappingBasenameToIds(names_ids_and_boxes, output_mapping)
def main(argv):
if len(argv) > 1:
raise RuntimeError('Too many command-line arguments.')
tf.logging.set_verbosity(tf.logging.INFO)
# Read list of index images from dataset file.
tf.logging.info('Reading list of index images from dataset file...')
_, index_list, _ = dataset.ReadDatasetFile(cmd_args.dataset_file_path)
num_images = len(index_list)
tf.logging.info('done! Found %d images', num_images)
# Parse DelfConfig proto.
config = delf_config_pb2.DelfConfig()
with tf.gfile.GFile(cmd_args.delf_config_path, 'r') as f:
text_format.Merge(f.read(), config)
# Create output directories if necessary.
if not os.path.exists(cmd_args.output_features_dir):
os.makedirs(cmd_args.output_features_dir)
if not os.path.exists(cmd_args.output_boxes_dir):
os.makedirs(cmd_args.output_boxes_dir)
index_names_ids_and_boxes = []
with tf.Graph().as_default():
with tf.Session() as sess:
# Initialize variables, construct detector and DELF extractor.
init_op = tf.global_variables_initializer()
sess.run(init_op)
detector_fn = extract_boxes.MakeDetector(
sess, cmd_args.detector_model_dir, import_scope='detector')
delf_extractor_fn = extract_features.MakeExtractor(
sess, config, import_scope='extractor_delf')
start = time.clock()
for i in range(num_images):
if i == 0:
print(
'Starting to extract features/boxes from index images...'
)
elif i % _STATUS_CHECK_ITERATIONS == 0:
elapsed = (time.clock() - start)
print('Processing index image %d out of %d, last %d '
'images took %f seconds' %
(i, num_images, _STATUS_CHECK_ITERATIONS, elapsed))
start = time.clock()
index_image_name = index_list[i]
input_image_filename = os.path.join(
cmd_args.images_dir, index_image_name + _IMAGE_EXTENSION)
output_feature_filename_whole_image = os.path.join(
cmd_args.output_features_dir,
index_image_name + _DELF_EXTENSION)
output_box_filename = os.path.join(
cmd_args.output_boxes_dir,
index_image_name + _BOX_EXTENSION)
pil_im = _PilLoader(input_image_filename)
width, height = pil_im.size
# Extract and save boxes.
if tf.gfile.Exists(output_box_filename):
tf.logging.info('Skipping box computation for %s',
index_image_name)
(boxes_out, scores_out, class_indices_out
) = box_io.ReadFromFile(output_box_filename)
else:
(boxes_out, scores_out, class_indices_out) = detector_fn(
np.expand_dims(pil_im, 0))
# Using only one image per batch.
boxes_out = boxes_out[0]
scores_out = scores_out[0]
class_indices_out = class_indices_out[0]
box_io.WriteToFile(output_box_filename, boxes_out,
scores_out, class_indices_out)
# Select boxes with scores greater than threshold. Those will be the
# ones with extracted DELF features (besides the whole image, whose DELF
# features are extracted in all cases).
num_delf_files = 1
selected_boxes = []
for box_ind, box in enumerate(boxes_out):
if scores_out[box_ind] >= cmd_args.detector_thresh:
selected_boxes.append(box)
num_delf_files += len(selected_boxes)
# Extract and save DELF features.
for delf_file_ind in range(num_delf_files):
if delf_file_ind == 0:
index_box_name = index_image_name
output_feature_filename = output_feature_filename_whole_image
else:
index_box_name = index_image_name + '_' + str(
delf_file_ind - 1)
output_feature_filename = os.path.join(
cmd_args.output_features_dir,
index_box_name + _DELF_EXTENSION)
index_names_ids_and_boxes.append(
[index_box_name, i, delf_file_ind - 1])
if tf.gfile.Exists(output_feature_filename):
tf.logging.info('Skipping DELF computation for %s',
index_box_name)
continue
if delf_file_ind >= 1:
bbox_for_cropping = selected_boxes[delf_file_ind - 1]
bbox_for_cropping_pil_convention = [
int(math.floor(bbox_for_cropping[1] * width)),
int(math.floor(bbox_for_cropping[0] * height)),
int(math.ceil(bbox_for_cropping[3] * width)),
int(math.ceil(bbox_for_cropping[2] * height))
]
pil_cropped_im = pil_im.crop(
bbox_for_cropping_pil_convention)
im = np.array(pil_cropped_im)
else:
im = np.array(pil_im)
(locations_out, descriptors_out, feature_scales_out,
attention_out) = delf_extractor_fn(im)
feature_io.WriteToFile(output_feature_filename,
locations_out, feature_scales_out,
descriptors_out, attention_out)
# Save mapping from output DELF name to index image id and box id.
_WriteMappingBasenameToIds(index_names_ids_and_boxes,
cmd_args.output_index_mapping)
