pythonでtfrecordsデータのレコードを選んでtfrecordsデータを再構成したい

Question

前提・実現したいこと

tfrecordsデータから特定のレコードのみを選び、それらからなる新しいtfrecordsデータを作成したいと考えています。

例えば、元のtfrecordsデータのレコード数が10だった場合に、2番目と3番目と7番目のレコードからなる新たなtfrecordsデータを構成したいです。

発生している問題・エラーメッセージ

自分で調べた限りでは、tfrecordsデータを分割する方法は見つけることができました。

しかし、任意のレコードを選んで再構成する方法は見つけることができませんでした。

最後に

質問に不備等ございましたらご指摘いただけますと幸いです。
ご回答よろしくお願いいたします。

Answer 1

最適解ではないように思いますが、自己解決しました。
拙いプログラム失礼します。

import tensorflow as tf


path = "" # 元のTFRecordファイルのパス


num_list =[] # 取り出したいレコード番号のリスト

rec_num = len(list(tf.compat.v1.python_io.tf_record_iterator(path)))
dataset = [] # 新しいTFRecordを構成するデータのリスト
count = 0

# TFRecordファイルの読み込み
raw_dataset = tf.data.TFRecordDataset(path)
for raw_record in raw_dataset.take(rec_num): #レコード数だけ繰り返す
    print(f"count : {count}")
    if str(count) in num_list:
        example = tf.train.Example()
        example.ParseFromString(raw_record.numpy())
        dataset.append(example)
    count += 1
    
# TFRecordファイルの書き出し
filename = "" # 書き出し先ファイルパス
with tf.io.TFRecordWriter(filename) as writer:
    for i in dataset:
        writer.write(i.SerializeToString())