Как дополнить данные обнаружения объекта в формате TFRecord - PullRequest
Новая
       81

Как дополнить данные обнаружения объекта в формате TFRecord

0 голосов
/ 20 апреля 2020

У меня есть небольшой набор данных (500 изображений, 10000 ярлыков, 16 классов), который я использую для обучения модели yolov3 . Я получаю приемлемые результаты после обучения, однако хочу улучшить показатели достоверности и точность обнаружения и, конечно, переоснащение не вариант. У меня есть данные тренировки, хранящиеся в формате TFRecord. Я заинтересован в использовании библиотеки imgaug, так как она имеет хорошую документацию и поддержку для работы с ограничивающими рамками. Вот последовательность шагов, которые я использую для обработки набора данных: 

train_dataset = read_tfr(FLAGS.dataset, FLAGS.classes, get_feature_map(), 100)
train_dataset = train_dataset.shuffle(buffer_size=512)
train_dataset = train_dataset.batch(FLAGS.batch_size)
train_dataset = train_dataset.map(lambda x, y: (
    dataset.transform_images(x, FLAGS.size),
    dataset.transform_targets(y, anchors, anchor_masks, FLAGS.size)))
train_dataset = train_dataset.prefetch(
    buffer_size=tf.data.experimental.AUTOTUNE)

Я хочу добавить еще train_dataset = train_dataset.map(some_augmentation_pipeline), и я не знаю, как этого можно достичь.

Здесь Вот функции чтения: 

def read_example(example, feature_map, class_table, max_boxes, new_size=None):
    """
    Read single a single example from a TFRecord file.
    Args:
        example: nd tensor.
        feature_map: A dictionary of feature names mapped to tf.io objects.
        class_table: StaticHashTable object.
        max_boxes: Maximum number of boxes per image
        new_size: w, h new image size

    Returns:
        x_train, y_train
    """
    features = tf.io.parse_single_example(example, feature_map)
    x_train = tf.image.decode_png(features['image_data'], channels=3)
    if new_size:
        x_train = tf.image.resize(x_train, new_size)
    object_name = tf.sparse.to_dense(features['object_name'])
    label = tf.cast(class_table.lookup(object_name), tf.float32)
    y_train = tf.stack([tf.sparse.to_dense(features[feature])
                        for feature in ['x_min', 'y_min', 'x_max', 'y_max']] + [label], 1)
    padding = [[0, max_boxes - tf.shape(y_train)[0]], [0, 0]]
    y_train = tf.pad(y_train, padding)
    return x_train, y_train

def read_tfr(tf_record_file, classes_file, feature_map, max_boxes,
             classes_delimiter='\n', new_size=None):
    """
    Read and load dataset from TFRecord file.
    Args:
        tf_record_file: Path to TFRecord file.
        classes_file: file containing classes.
        feature_map: A dictionary of feature names mapped to tf.io objects.
        max_boxes: Maximum number of boxes per image.
        classes_delimiter: delimiter in classes_file.
        new_size: w, h new image size

    Returns:
        MapDataset object.
    """
    text_init = tf.lookup.TextFileInitializer(
        classes_file, tf.string, 0, tf.int64, -1, delimiter=classes_delimiter)
    class_table = tf.lookup.StaticHashTable(text_init, -1)
    files = tf.data.Dataset.list_files(tf_record_file)
    dataset = files.flat_map(tf.data.TFRecordDataset)
    return dataset.map(
        lambda x: read_example(x, feature_map, class_table, max_boxes, new_size))

Я думаю, может быть, я предварительно увеличу изображения и сохраню их в одном и том же TFRecord вместе с их настроенными метками / ограничивающими рамками, но я не думаю, что это лучше способ сделать это ... что вы думаете?

...