def test_should_use_color_map_keys_as_channels_by_default(self):
with TemporaryDirectory() as path:
tfrecord_filename = os.path.join(path, 'data.tfrecord')
get_logger().debug('writing to test tfrecord_filename: %s',
tfrecord_filename)
write_examples_to_tfrecord(
tfrecord_filename,
[dict_to_example({'image': encode_png([[COLOR_1, COLOR_2]])})])
class_weights_map = calculate_median_class_weights_for_tfrecord_paths_and_color_map(
[tfrecord_filename], 'image', {
'color1': COLOR_1,
'color2': COLOR_2
})
assert set(class_weights_map.keys()) == {'color1', 'color2'}
def test_should_calculate_median_class_weights_for_multiple_image_and_multiple_images(
self):
with TemporaryDirectory() as path:
tfrecord_filename = os.path.join(path, 'data.tfrecord')
get_logger().debug('writing to test tfrecord_filename: %s',
tfrecord_filename)
write_examples_to_tfrecord(tfrecord_filename, [
dict_to_example(
{'image': encode_png([[COLOR_0, COLOR_1, COLOR_2]])}),
dict_to_example(
{'image': encode_png([[COLOR_1, COLOR_2, COLOR_3]])})
])
class_weights = calculate_median_class_weights_for_tfrecord_paths_and_colors(
[tfrecord_filename], 'image', [COLOR_1, COLOR_2, COLOR_3])
assert class_weights == [0.25, 0.25, 0.5]
def test_should_calculate_median_class_weights_for_single_image_and_single_color(
self):
with TemporaryDirectory() as path:
tfrecord_filename = os.path.join(path, 'data.tfrecord')
get_logger().debug('writing to test tfrecord_filename: %s',
tfrecord_filename)
write_examples_to_tfrecord(
tfrecord_filename,
[dict_to_example({'image': encode_png([[COLOR_1, COLOR_2]])})])
class_weights_map = calculate_median_class_weights_for_tfrecord_paths_and_color_map(
[tfrecord_filename],
'image', {
'color1': COLOR_1,
'color2': COLOR_2,
'color3': COLOR_3
},
channels=['color1', 'color2'])
assert class_weights_map == {'color1': 0.5, 'color2': 0.5}
def test_should_include_unknown_class_if_enabled(self):
with TemporaryDirectory() as path:
tfrecord_filename = os.path.join(path, 'data.tfrecord')
get_logger().debug('writing to test tfrecord_filename: %s',
tfrecord_filename)
write_examples_to_tfrecord(tfrecord_filename, [
dict_to_example({
'image':
encode_png([[COLOR_0, COLOR_1, COLOR_2, COLOR_3]])
})
])
class_weights_map = calculate_median_class_weights_for_tfrecord_paths_and_color_map(
[tfrecord_filename],
'image', {
'color1': COLOR_1,
'color2': COLOR_2
},
use_unknown_class=True,
unknown_class_label='unknown')
assert set(
class_weights_map.keys()) == {'color1', 'color2', 'unknown'}