def test_raises_error_with_no_input_paths(self):
input_reader_text_proto = """
shuffle: false
num_readers: 1
load_instance_masks: true
"""
input_reader_proto = input_reader_pb2.InputReader()
text_format.Merge(input_reader_text_proto, input_reader_proto)
with self.assertRaises(ValueError):
dataset_builder.build(input_reader_proto)
def test_build_coco_stuff(cfg):
dataset = dataset_builder.build(cfg)
assert isinstance(dataset, Dataset)
for datum in dataset:
seg = datum['gt-seg'].numpy()
assert len(seg[seg == -1]) == 0
def test_build_tf_record_input_reader(self):
tf_record_path = self.create_tf_record()
input_reader_text_proto = """
shuffle: false
num_readers: 1
tf_record_input_reader {{
input_path: '{0}'
}}
""".format(tf_record_path)
input_reader_proto = input_reader_pb2.InputReader()
text_format.Merge(input_reader_text_proto, input_reader_proto)
tensor_dict = dataset_util.make_initializable_iterator(
dataset_builder.build(input_reader_proto)).get_next()
sv = tf.train.Supervisor(logdir=self.get_temp_dir())
with sv.prepare_or_wait_for_session() as sess:
sv.start_queue_runners(sess)
output_dict = sess.run(tensor_dict)
self.assertTrue(fields.InputDataFields.groundtruth_instance_masks
not in output_dict)
self.assertEquals((4, 5, 3),
output_dict[fields.InputDataFields.image].shape)
self.assertEquals(
[2], output_dict[fields.InputDataFields.groundtruth_classes])
self.assertEquals(
(1, 4),
output_dict[fields.InputDataFields.groundtruth_boxes].shape)
self.assertAllEqual(
[0.0, 0.0, 1.0, 1.0],
output_dict[fields.InputDataFields.groundtruth_boxes][0])
def build_single_sampler(type, cfg):
dataset = dataset_builder.build(cfg['dataset'])
single_sampler_class = load_single_sampler_class(type)
sampler = single_sampler_class.build(dataset, cfg)
return sampler, dataset
def extract_images(gpus, model_config, data_config, trained_checkpoint,
pad_to_shape, processor_type, annot_type, is_debug,
export_folder, **kwargs):
os.environ["CUDA_VISIBLE_DEVICES"] = gpus
pipeline_config = read_config(model_config, data_config)
if pad_to_shape is not None and isinstance(pad_to_shape, str):
pad_to_shape = [
int(dim) if dim != '-1' else None
for dim in pad_to_shape.split(',')
]
input_reader = pipeline_config.input_reader
input_reader.shuffle = False
if len(input_reader.tf_record_input_reader) > 1:
input_reader.tf_record_input_reader.pop()
print("REMOVED INPUT READER:\n", input_reader)
ignore_label = input_reader.ignore_label
num_classes, segmentation_model = model_builder.build(
pipeline_config.model, is_training=False, ignore_label=ignore_label)
with tf.device("cpu:0"):
dataset = dataset_builder.build(input_reader, 1)
num_gpu = len(gpus.split(","))
num_examples = sum(
[r.num_examples for r in input_reader.tf_record_input_reader])
run_inference_graph(segmentation_model, trained_checkpoint, dataset,
num_examples, ignore_label, pad_to_shape, num_classes,
processor_type, annot_type, num_gpu, export_folder,
**kwargs)
def test_reid_writing():
gallery_dataset = dataset_builder.build(gallery_cfg)
query_dataset = dataset_builder.build(gallery_cfg)
gallery_dataloader = torch.utils.data.DataLoader(gallery_dataset,
batch_size=10,
num_workers=0)
model_cfg = model_cfg_builder.build(cfg)
model = model_builder.build(model_cfg[0])
model = DataParallel(model)
output_file = tempfile.mktemp()
objgraph.show_growth()
write_to_h5(gallery_dataloader, model, output_file)
objgraph.show_most_common_types()
objgraph.show_growth()
time.sleep(1)
query_dataloader = torch.utils.data.DataLoader(query_dataset,
batch_size=10,
num_workers=0)
output_file = tempfile.mktemp()
write_to_h5(query_dataloader, model, output_file)
def test_transform():
dataset = dataset_builder.build(mpii_config)
joint_info = dataset.info['joint_info']
for _ in range(10):
idx = np.random.randint(0, len(dataset))
#idx = 15707
data = dataset[idx]
img = data['img']
coords = data['coords']
#print("after to tensor", img)
visualize(img, coords, joint_info.stick_figure_edges)
def run_experiment(gpus, model_config, data_config, trained_checkpoint,
pad_to_shape, processor_type, annot_type, is_debug,
**kwargs):
had_error = None
try:
os.environ["CUDA_VISIBLE_DEVICES"] = gpus
print_buffer = StringIO()
if not is_debug:
sys.stdout = print_buffer
sys.stderr = print_buffer
pipeline_config = read_config(model_config, data_config)
if pad_to_shape is not None and isinstance(pad_to_shape, str):
pad_to_shape = [
int(dim) if dim != '-1' else None
for dim in pad_to_shape.split(',')
]
input_reader = pipeline_config.input_reader
input_reader.shuffle = True
ignore_label = input_reader.ignore_label
num_classes, segmentation_model = model_builder.build(
pipeline_config.model,
is_training=False,
ignore_label=ignore_label)
with tf.device("cpu:0"):
dataset = dataset_builder.build(input_reader, 1)
num_gpu = len(gpus.split(","))
num_examples = sum(
[r.num_examples for r in input_reader.tf_record_input_reader])
result = run_inference_graph(segmentation_model, trained_checkpoint,
dataset, num_examples, ignore_label,
pad_to_shape, num_classes, processor_type,
annot_type, num_gpu, **kwargs)
had_error = False
except Exception as ex:
if is_debug:
raise ex
print(traceback.format_exc())
had_error = True
result = None
sys.stdout = sys.__stdout__
sys.stderr = sys.__stderr__
tf.reset_default_graph()
return print_buffer, result, had_error
def main(_):
#test_plots()
eval_dir = FLAGS.eval_dir
output_directory = os.path.join(eval_dir, "inf")
suff = ""
if FLAGS.global_mean:
suff = "_G"
else:
suff = "_L"
if FLAGS.global_cov:
suff += "G"
else:
suff += "L"
dist_dir = os.path.join(eval_dir, "class_dist" + suff)
min_dir = os.path.join(eval_dir, "min" + suff)
hist_dir = os.path.join(eval_dir, "hist" + suff)
dump_dir = os.path.join(eval_dir, "dump" + suff)
tf.gfile.MakeDirs(output_directory)
tf.gfile.MakeDirs(min_dir)
tf.gfile.MakeDirs(dist_dir)
tf.gfile.MakeDirs(hist_dir)
pipeline_config = read_config(FLAGS.model_config, FLAGS.data_config)
pad_to_shape = None
if FLAGS.input_shape:
input_shape = [
int(dim) if dim != '-1' else None
for dim in FLAGS.input_shape.split(',')]
else:
raise ValueError('Must supply `input_shape`')
if FLAGS.pad_to_shape:
pad_to_shape = [
int(dim) if dim != '-1' else None
for dim in FLAGS.pad_to_shape.split(',')]
label_map = (CITYSCAPES_LABEL_IDS
if FLAGS.label_ids else CITYSCAPES_LABEL_COLORS)
input_reader = pipeline_config.input_reader
input_reader.shuffle = True
ignore_label = input_reader.ignore_label
num_classes, segmentation_model = model_builder.build(
pipeline_config.model, is_training=False, ignore_label=ignore_label)
dataset = dataset_builder.build(input_reader, 1)
run_inference_graph(segmentation_model, FLAGS.trained_checkpoint,
dataset, input_reader.num_examples, ignore_label, input_shape, pad_to_shape,
label_map, output_directory, num_classes, eval_dir, min_dir, dist_dir, hist_dir,
dump_dir)
def test_pose_writing():
pose_dataset = dataset_builder.build(pose_dataset_cfg)
pose_dataloader = torch.utils.data.DataLoader(pose_dataset,
batch_size=10,
num_workers=0)
model_cfg = evaluation_model_builder.build(pose_model_cfg)
model = model_builder.build(model_cfg[0])
model = DataParallel(model)
filename = "tests/pose.h5"
if os.path.isfile(filename):
os.remove(filename)
print("deleted old {}".format(filename))
write_to_h5(pose_dataloader, model, "tests/pose.h5", ["pose"])
def test_pose_writing():
pose_dataset = dataset_builder.build(pose_dataset_cfg)
pose_dataloader = torch.utils.data.DataLoader(pose_dataset,
batch_size=10,
num_workers=0)
model_cfg = model_cfg_builder.build(pose_model_cfg)[0]
dataset_info = pose_dataset.info
new_cfg = dataset_info.copy()
new_cfg.update(model_cfg)
print(new_cfg)
model = model_builder.build(new_cfg)
model = DataParallel(model)
output_file = './tests/pose.h5'
write_to_h5(pose_dataloader, model, output_file, ['emb', 'pose'])
def test_market_attribute_dataset():
market_attribute_cfg = {
"data_dir": Config.MARKET_ATTRIBUTE,
"split": 'train',
'name': 'market1501_attribute'
}
data = dataset_builder.build(market_attribute_cfg)
assert data[0]['hat'] == 0
assert data[174]['hat'] == 1
assert data[0]['upcolor'] == 2
assert data[0]['downcolor'] == 6
for idx, d in enumerate(data):
assert (d['downcolor']) != 9, idx
def test_build_tf_record_input_reader_with_batch_size_two(self):
tf_record_path = self.create_tf_record()
input_reader_text_proto = """
shuffle: false
num_readers: 1
tf_record_input_reader {{
input_path: '{0}'
}}
""".format(tf_record_path)
input_reader_proto = input_reader_pb2.InputReader()
text_format.Merge(input_reader_text_proto, input_reader_proto)
def one_hot_class_encoding_fn(tensor_dict):
tensor_dict[
fields.InputDataFields.groundtruth_classes] = tf.one_hot(
tensor_dict[fields.InputDataFields.groundtruth_classes] -
1,
depth=3)
return tensor_dict
tensor_dict = dataset_util.make_initializable_iterator(
dataset_builder.build(
input_reader_proto,
transform_input_data_fn=one_hot_class_encoding_fn,
batch_size=2,
max_num_boxes=2,
num_classes=3,
spatial_image_shape=[4, 5])).get_next()
sv = tf.train.Supervisor(logdir=self.get_temp_dir())
with sv.prepare_or_wait_for_session() as sess:
sv.start_queue_runners(sess)
output_dict = sess.run(tensor_dict)
self.assertAllEqual([2, 4, 5, 3],
output_dict[fields.InputDataFields.image].shape)
self.assertAllEqual(
[2, 2, 3],
output_dict[fields.InputDataFields.groundtruth_classes].shape)
self.assertAllEqual(
[2, 2, 4],
output_dict[fields.InputDataFields.groundtruth_boxes].shape)
self.assertAllEqual(
[[[0.0, 0.0, 1.0, 1.0], [0.0, 0.0, 0.0, 0.0]],
[[0.0, 0.0, 1.0, 1.0], [0.0, 0.0, 0.0, 0.0]]],
output_dict[fields.InputDataFields.groundtruth_boxes])
def main(_):
assert FLAGS.output_dir, '`output_dir` missing.'
output_directory = FLAGS.output_dir
tf.gfile.MakeDirs(output_directory)
pipeline_config = read_config(FLAGS.model_config, FLAGS.data_config)
pad_to_shape = None
if FLAGS.input_shape:
input_shape = [
int(dim) if dim != '-1' else None
for dim in FLAGS.input_shape.split(',')
]
else:
raise ValueError('Must supply `input_shape`')
patch_size = None
if FLAGS.patch_size:
patch_size = [int(dim) for dim in FLAGS.patch_size.split(',')]
assert len(patch_size) == 2, "patch size must be h,w"
if FLAGS.pad_to_shape:
pad_to_shape = [
int(dim) if dim != '-1' else None
for dim in FLAGS.pad_to_shape.split(',')
]
label_map = (CITYSCAPES_LABEL_IDS
if FLAGS.label_ids else CITYSCAPES_LABEL_COLORS)
ignore_label = pipeline_config.input_reader.ignore_label
num_classes, segmentation_model = model_builder.build(
pipeline_config.model, ignore_label=ignore_label, is_training=False)
#input_reader = pipeline_config.eval_input_reader
input_reader = pipeline_config.input_reader
input_reader.shuffle = False
input_dict = dataset_builder.build(input_reader, epoch)
num_examples = sum(
[r.num_examples for r in input_reader.tf_record_input_reader])
iters = min(num_examples, FLAGS.max_iters)
run_inference_graph(segmentation_model, FLAGS.trained_checkpoint,
input_dict, iters, input_shape, pad_to_shape,
label_map, output_directory, num_classes, patch_size)
def build(cfg):
if isinstance(cfg, list):
dataloaders = []
for c in cfg:
dataloaders.extend(build(c))
return dataloaders
else:
num_workers = cfg.get('num_workers', 4)
pin_memory = cfg.get('pin_memory', True)
print("num_workers", num_workers)
if num_workers == -1:
import multiprocessing
# one for main thread
num_workers = multiprocessing.cpu_count() - 1
print("num_workers", num_workers)
if 'sampler' in cfg:
sampler, dataset = sampler_builder.build(cfg['sampler'])
if isinstance(dataset, MultiDataset):
collate_fn = build_collate_fn(dataset.header)
else:
collate_fn = default_collate
dataloader = torch.utils.data.DataLoader(dataset,
batch_sampler=sampler,
collate_fn=collate_fn,
num_workers=num_workers,
pin_memory=pin_memory)
else:
# use default
dataset = dataset_builder.build(cfg['dataset'])
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=cfg.get(
'batch_size', 10),
num_workers=num_workers,
pin_memory=pin_memory)
return [dataloader]
def test_build_tf_record_input_reader_and_load_instance_masks(self):
tf_record_path = self.create_tf_record()
input_reader_text_proto = """
shuffle: false
num_readers: 1
load_instance_masks: true
tf_record_input_reader {{
input_path: '{0}'
}}
""".format(tf_record_path)
input_reader_proto = input_reader_pb2.InputReader()
text_format.Merge(input_reader_text_proto, input_reader_proto)
tensor_dict = dataset_util.make_initializable_iterator(
dataset_builder.build(input_reader_proto, batch_size=1)).get_next()
sv = tf.train.Supervisor(logdir=self.get_temp_dir())
with sv.prepare_or_wait_for_session() as sess:
sv.start_queue_runners(sess)
output_dict = sess.run(tensor_dict)
self.assertAllEqual(
(1, 1, 4, 5),
output_dict[fields.InputDataFields.groundtruth_instance_masks].shape)
def test_duke_attribute_dataset():
duke_attribute_cfg = {
"data_dir": Config.DUKE_ATTRIBUTE,
"split": 'train',
'name': 'duke_mtmc_attribute'
}
data = dataset_builder.build(duke_attribute_cfg)
assert data[0]['hat'] == 0
assert data[4]['hat'] == 1
assert data[7]['upcolor'] == 5
assert data[336]['downcolor'] == 5
for idx, d in enumerate(data):
assert (d['gender']) < 2, idx
assert (d['top']) < 2, idx
assert (d['boots']) < 2, idx
assert (d['hat']) < 2, idx
assert (d['backpack']) < 2, idx
assert (d['bag']) < 2, idx
assert (d['handbag']) < 2, idx
assert (d['shoes']) < 2, idx
assert (d['upcolor']) < 8, idx
assert (d['downcolor']) < 7, idx
def test_reid_datasets_builder():
dataset = dataset_builder.build(reid_dataset_cfgs)
assert isinstance(dataset, ConcatReidDataset)
def test_reid_dataset_builder():
dataset = dataset_builder.build(reid_dataset_cfg)
assert isinstance(dataset, ReidDataset)
def _eval_input_fn(params=None):
"""Returns `features` and `labels` tensor dictionaries for evaluation.
Args:
params: Parameter dictionary passed from the estimator.
Returns:
features: Dictionary of feature tensors.
features[fields.InputDataFields.image] is a [1, H, W, C] float32 tensor
with preprocessed images.
features[HASH_KEY] is a [1] int32 tensor representing unique
identifiers for the images.
features[fields.InputDataFields.true_image_shape] is a [1, 3]
int32 tensor representing the true image shapes, as preprocessed
images could be padded.
features[fields.InputDataFields.original_image] is a [1, H', W', C]
float32 tensor with the original image.
labels: Dictionary of groundtruth tensors.
labels[fields.InputDataFields.groundtruth_boxes] is a [1, num_boxes, 4]
float32 tensor containing the corners of the groundtruth boxes.
labels[fields.InputDataFields.groundtruth_classes] is a
[num_boxes, num_classes] float32 one-hot tensor of classes.
labels[fields.InputDataFields.groundtruth_area] is a [1, num_boxes]
float32 tensor containing object areas.
labels[fields.InputDataFields.groundtruth_is_crowd] is a [1, num_boxes]
bool tensor indicating if the boxes enclose a crowd.
labels[fields.InputDataFields.groundtruth_difficult] is a [1, num_boxes]
int32 tensor indicating if the boxes represent difficult instances.
-- Optional --
labels[fields.InputDataFields.groundtruth_instance_masks] is a
[1, num_boxes, H, W] float32 tensor containing only binary values,
which represent instance masks for objects.
Raises:
TypeError: if the `eval_config` or `eval_input_config` are not of the
correct type.
"""
del params
if not isinstance(eval_config, eval_pb2.EvalConfig):
raise TypeError('For eval mode, the `eval_config` must be a '
'train_pb2.EvalConfig.')
if not isinstance(eval_input_config, input_reader_pb2.InputReader):
raise TypeError('The `eval_input_config` must be a '
'input_reader_pb2.InputReader.')
if not isinstance(model_config, model_pb2.DetectionModel):
raise TypeError('The `model_config` must be a '
'model_pb2.DetectionModel.')
num_classes = config_util.get_number_of_classes(model_config)
model = model_builder.build(model_config, is_training=False)
image_resizer_config = config_util.get_image_resizer_config(
model_config)
image_resizer_fn = image_resizer_builder.build(image_resizer_config)
transform_data_fn = functools.partial(
transform_input_data,
model_preprocess_fn=model.preprocess,
image_resizer_fn=image_resizer_fn,
num_classes=num_classes,
data_augmentation_fn=None,
retain_original_image=True)
dataset = dataset_builder.build(
eval_input_config, transform_input_data_fn=transform_data_fn)
input_dict = dataset_util.make_initializable_iterator(
dataset).get_next()
hash_from_source_id = tf.string_to_hash_bucket_fast(
input_dict[fields.InputDataFields.source_id], HASH_BINS)
features = {
fields.InputDataFields.image:
input_dict[fields.InputDataFields.image],
fields.InputDataFields.original_image:
input_dict[fields.InputDataFields.original_image],
HASH_KEY:
tf.cast(hash_from_source_id, tf.int32),
fields.InputDataFields.true_image_shape:
input_dict[fields.InputDataFields.true_image_shape]
}
labels = {
fields.InputDataFields.groundtruth_boxes:
input_dict[fields.InputDataFields.groundtruth_boxes],
fields.InputDataFields.groundtruth_classes:
input_dict[fields.InputDataFields.groundtruth_classes],
fields.InputDataFields.groundtruth_area:
input_dict[fields.InputDataFields.groundtruth_area],
fields.InputDataFields.groundtruth_is_crowd:
input_dict[fields.InputDataFields.groundtruth_is_crowd],
fields.InputDataFields.groundtruth_difficult:
tf.cast(input_dict[fields.InputDataFields.groundtruth_difficult],
tf.int32)
}
if fields.InputDataFields.groundtruth_instance_masks in input_dict:
labels[fields.InputDataFields.
groundtruth_instance_masks] = input_dict[
fields.InputDataFields.groundtruth_instance_masks]
# Add a batch dimension to the tensors.
features = {
key: tf.expand_dims(features[key], axis=0)
for key, feature in features.items()
}
labels = {
key: tf.expand_dims(labels[key], axis=0)
for key, label in labels.items()
}
return features, labels
def test_pose_dataset():
dataset = dataset_builder.build(mpii_config)
for data in dataset:
assert data['img'].shape == (3, 256, 256)
print(data['coords'])
break
def main(_):
#test_plots()
eval_dir = FLAGS.eval_dir
output_directory = os.path.join(eval_dir, "inf")
suff = ""
if FLAGS.global_mean:
suff = "_G"
else:
suff = "_L"
if FLAGS.global_cov:
suff += "G"
else:
suff += "L"
dist_dir = os.path.join(eval_dir, "class_dist" + suff)
min_dir = os.path.join(eval_dir, "min" + suff)
hist_dir = os.path.join(eval_dir, "hist" + suff)
dump_dir = os.path.join(eval_dir, "dump" + suff)
tf.gfile.MakeDirs(output_directory)
tf.gfile.MakeDirs(min_dir)
tf.gfile.MakeDirs(dist_dir)
tf.gfile.MakeDirs(hist_dir)
pipeline_config = pipeline_pb2.PipelineConfig()
with tf.gfile.GFile(FLAGS.config_path, 'r') as f:
text_format.Merge(f.read(), pipeline_config)
pad_to_shape = None
if FLAGS.input_shape:
input_shape = [
int(dim) if dim != '-1' else None
for dim in FLAGS.input_shape.split(',')
]
else:
raise ValueError('Must supply `input_shape`')
if FLAGS.pad_to_shape:
pad_to_shape = [
int(dim) if dim != '-1' else None
for dim in FLAGS.pad_to_shape.split(',')
]
label_map = (CITYSCAPES_LABEL_IDS
if FLAGS.label_ids else CITYSCAPES_LABEL_COLORS)
num_classes, segmentation_model = model_builder.build(
pipeline_config.model, is_training=False)
if FLAGS.do_ood:
if FLAGS.write_out or FLAGS.use_train:
input_reader = pipeline_config.ood_train_input_reader
else:
input_reader = pipeline_config.ood_eval_input_reader
else:
if FLAGS.use_train:
input_reader = pipeline_config.train_input_reader
else:
input_reader = pipeline_config.eval_input_reader
input_reader.shuffle = True
input_reader.num_epochs = 1
input_reader.num_examples = min(1500, input_reader.num_examples)
input_dict = dataset_builder.build(input_reader)
ignore_label = pipeline_config.ood_config.ignore_label
run_inference_graph(segmentation_model, FLAGS.trained_checkpoint,
input_dict, input_reader.num_examples, ignore_label,
input_shape, pad_to_shape, label_map, output_directory,
num_classes, eval_dir, min_dir, dist_dir, hist_dir,
dump_dir)
def get_next(config):
return dataset_util.make_initializable_iterator(
dataset_builder.build(config)).get_next()
def test_pose_dataset():
dataset = dataset_builder.build(mpii_config)
assert isinstance(dataset, PoseDataset)
def test_mpii_dataset():
dataset = dataset_builder.build(mpii_config)
assert isinstance(dataset, Mpii)
def _train_input_fn(params=None):
"""Returns `features` and `labels` tensor dictionaries for training.
Args:
params: Parameter dictionary passed from the estimator.
Returns:
features: Dictionary of feature tensors.
features[fields.InputDataFields.image] is a [batch_size, H, W, C]
float32 tensor with preprocessed images.
features[HASH_KEY] is a [batch_size] int32 tensor representing unique
identifiers for the images.
features[fields.InputDataFields.true_image_shape] is a [batch_size, 3]
int32 tensor representing the true image shapes, as preprocessed
images could be padded.
labels: Dictionary of groundtruth tensors.
labels[fields.InputDataFields.num_groundtruth_boxes] is a [batch_size]
int32 tensor indicating the number of groundtruth boxes.
labels[fields.InputDataFields.groundtruth_boxes] is a
[batch_size, num_boxes, 4] float32 tensor containing the corners of
the groundtruth boxes.
labels[fields.InputDataFields.groundtruth_classes] is a
[batch_size, num_boxes, num_classes] float32 one-hot tensor of
classes.
labels[fields.InputDataFields.groundtruth_weights] is a
[batch_size, num_boxes] float32 tensor containing groundtruth weights
for the boxes.
-- Optional --
labels[fields.InputDataFields.groundtruth_instance_masks] is a
[batch_size, num_boxes, H, W] float32 tensor containing only binary
values, which represent instance masks for objects.
labels[fields.InputDataFields.groundtruth_keypoints] is a
[batch_size, num_boxes, num_keypoints, 2] float32 tensor containing
keypoints for each box.
Raises:
TypeError: if the `train_config` or `train_input_config` are not of the
correct type.
"""
if not isinstance(train_config, train_pb2.TrainConfig):
raise TypeError('For training mode, the `train_config` must be a '
'train_pb2.TrainConfig.')
if not isinstance(train_input_config, input_reader_pb2.InputReader):
raise TypeError('The `train_input_config` must be a '
'input_reader_pb2.InputReader.')
if not isinstance(model_config, model_pb2.DetectionModel):
raise TypeError('The `model_config` must be a '
'model_pb2.DetectionModel.')
data_augmentation_options = [
preprocessor_builder.build(step)
for step in train_config.data_augmentation_options
]
data_augmentation_fn = functools.partial(
augment_input_data,
data_augmentation_options=data_augmentation_options)
model = model_builder.build(model_config, is_training=True)
image_resizer_config = config_util.get_image_resizer_config(
model_config)
image_resizer_fn = image_resizer_builder.build(image_resizer_config)
transform_data_fn = functools.partial(
transform_input_data,
model_preprocess_fn=model.preprocess,
image_resizer_fn=image_resizer_fn,
num_classes=config_util.get_number_of_classes(model_config),
data_augmentation_fn=data_augmentation_fn)
dataset = dataset_builder.build(
train_input_config,
transform_input_data_fn=transform_data_fn,
batch_size=params['batch_size']
if params else train_config.batch_size,
max_num_boxes=train_config.max_number_of_boxes,
num_classes=config_util.get_number_of_classes(model_config),
spatial_image_shape=config_util.get_spatial_image_size(
image_resizer_config))
tensor_dict = dataset_util.make_initializable_iterator(
dataset).get_next()
hash_from_source_id = tf.string_to_hash_bucket_fast(
tensor_dict[fields.InputDataFields.source_id], HASH_BINS)
features = {
fields.InputDataFields.image:
tensor_dict[fields.InputDataFields.image],
HASH_KEY:
tf.cast(hash_from_source_id, tf.int32),
fields.InputDataFields.true_image_shape:
tensor_dict[fields.InputDataFields.true_image_shape]
}
labels = {
fields.InputDataFields.num_groundtruth_boxes:
tensor_dict[fields.InputDataFields.num_groundtruth_boxes],
fields.InputDataFields.groundtruth_boxes:
tensor_dict[fields.InputDataFields.groundtruth_boxes],
fields.InputDataFields.groundtruth_classes:
tensor_dict[fields.InputDataFields.groundtruth_classes],
fields.InputDataFields.groundtruth_weights:
tensor_dict[fields.InputDataFields.groundtruth_weights]
}
if fields.InputDataFields.groundtruth_keypoints in tensor_dict:
labels[fields.InputDataFields.groundtruth_keypoints] = tensor_dict[
fields.InputDataFields.groundtruth_keypoints]
if fields.InputDataFields.groundtruth_instance_masks in tensor_dict:
labels[fields.InputDataFields.
groundtruth_instance_masks] = tensor_dict[
fields.InputDataFields.groundtruth_instance_masks]
return features, labels
