2edd1fe70a43b5b646662ef6603f579275141fae,examples/utils/tf_records_generation.py,,main,#,96

Before Change

        shutil.copy(tile_dir, train_dir)

    // for tile in test_df["filename"]:
    for tile in list(tf_test[:,1]):
        tile_dir = op.join(tiles_dir, tile)
        shutil.copy(tile_dir, test_dir)
    ////// for train

After Change

    column_name = ["filename", "width", "height", "class", "xmin", "ymin", "xmax", "ymax"]
    df = pd.DataFrame(tf_tiles_info, columns=column_name)
    //shuffle the dataframe
    df = df.sample(frac=1)
    train_df = df[:split_index]
    test_df = df[split_index:]
    print("You have {} training tiles and {} test tiles ready".format(
    len(set(train_df["filename"])), len(set(test_df["filename"]))))
    // train_df.to_csv("train_df.csv")
    // test_df.to_csv("test_df.csv")

    tiles_dir = op.join(os.getcwd(), "tiles")
    train_dir = op.join(os.getcwd(), "images", "train")
    test_dir = op.join(os.getcwd(), "images", "test")

    if not op.isdir(train_dir):
        makedirs(train_dir)
    if not op.isdir(test_dir):
        makedirs(test_dir)

    for tile in train_df["filename"]:
        tile_dir = op.join(tiles_dir, tile)
        shutil.copy(tile_dir, train_dir)

    for tile in test_df["filename"]:
        tile_dir = op.join(tiles_dir, tile)
        shutil.copy(tile_dir, test_dir)
    ////// for train
    writer = tf.python_io.TFRecordWriter(FLAGS.train_rd_path)
    grouped = split(train_df, "filename")

    for group in grouped:
        tf_example = create_tf_example(group, train_dir)
        writer.write(tf_example.SerializeToString())
    writer.close()
    output_train= op.join(os.getcwd(),FLAGS.train_rd_path)
    print("Successfully created the TFRecords: {}".format(output_train))

Instances