def get_num_lightings(yaml_dict):
'''
Returns the number of non-blank lighting values.
'''
fg_path = yaml_dict['datasets']['training']['fg_path']
dataset = MemmapDataset(os.path.join(data_path, fg_path))
lighting_labels = dataset.tensors[1][:, 4]
assert_equal(lighting_labels[0], -1)
assert_array_compare(numpy.greater_equal, lighting_labels[1:], 0)
assert_array_compare(numpy.less, lighting_labels[1:], 4)
num_valid_lighting_values = len(frozenset(lighting_labels[1:]))
assert_equal(num_valid_lighting_values, 4)
return num_valid_lighting_values
def __init__(self,
model,
foreground_dataset,
video_frame,
original_image_size,
output_dir=None):
'''
model: IdAndCameraDirModel
foreground_dataset: simplelearn.data.Dataset
A Dataset with NORB-style label vectors. Used to compute a mapping
from object ID to example images.
video_frame: numpy.ndarray
A frame from the input video stream. Used to initialize display pixel
buffer's size and dtype.
output_dir: string or None
The directory to save frame images to. If omitted, no frame images
will be saved.
'''
assert_is_instance(model, (IdAndCameraDirModel,
IdAndCameraDirModelConv,
VideoModel))
assert_equal(model.input_node.output_format.axes,
('b', '0', '1', 'c'))
assert_is_instance(foreground_dataset, Dataset)
assert_equal(len(video_frame.shape), 3)
assert_in(video_frame.shape[-1], (1, 3))
original_image_size = numpy.asarray(original_image_size)
assert_equal(len(original_image_size), 2)
assert_array_compare(numpy.greater, original_image_size, 0)
assert_is(output_dir, None) # TODO: support this argument
self.original_image_size = numpy.asarray(original_image_size)
self.label_to_id = NorbLabelToObjectIdConverter(
foreground_dataset.tensors[1])
def make_id_elev_azim_to_example(example_dataset):
assert_equal(example_dataset.formats[0].axes.index('c'), 3)
assert_equal(example_dataset.formats[0].axes.index('b'), 0)
assert_equal(example_dataset.formats[1].axes.index('b'), 0)
images = example_dataset.tensors[0][..., :3] # cut off alpha
labels = example_dataset.tensors[1]
assert_in(labels.shape[1], (5, 11))
# Arbitrarily restrict attention to images that use
# lighting setup 0.
row_mask = labels[:, 4] == 0
images = images[row_mask, :]
labels = labels[row_mask, :]
ids = self.label_to_id(labels)
ids_elevs_azims = numpy.hstack((ids[:, numpy.newaxis],
labels[:, 2:4]))
result = dict(safe_izip((tuple(t) for t in ids_elevs_azims),
images))
assert_equal(len(result), ids_elevs_azims.shape[0])
return result
self.id_elev_azim_to_example = \
make_id_elev_azim_to_example(foreground_dataset)
self.model = model
self.all_pixels = numpy.zeros([video_frame.shape[0],
video_frame.shape[1] * 2,
video_frame.shape[2]],
dtype=video_frame.dtype)
self.video_pixels = self.all_pixels[:, :video_frame.shape[1], :]
self.examples_pixels = self.all_pixels[:, video_frame.shape[1]:, :]
