在TensorFlow中構建目標檢測數據集可以使用tf.data.Dataset類。以下是一個示例從圖片路徑和標注文件構建目標檢測數據集的方法:
import tensorflow as tf
import numpy as np
# 讀取圖片路徑和標注文件
image_paths = ["image1.jpg", "image2.jpg", "image3.jpg"]
annotations = [{"bbox": [100, 100, 200, 200], "class": 0},
{"bbox": [150, 150, 250, 250], "class": 1},
{"bbox": [200, 200, 300, 300], "class": 0}]
# 構建數據集
def parse_function(image_path, annotation):
# 讀取圖片
image = tf.io.read_file(image_path)
image = tf.io.decode_jpeg(image, channels=3)
# 處理標注
bbox = annotation["bbox"]
class_id = annotation["class"]
# 返回圖片和標注
return image, bbox, class_id
# 創建數據集
dataset = tf.data.Dataset.from_tensor_slices((image_paths, annotations))
dataset = dataset.map(lambda image_path, annotation: tf.py_function(parse_function, [image_path, annotation], [tf.uint8, tf.float32, tf.int32]))
# 打亂數據集并設置batch size
dataset = dataset.shuffle(buffer_size=len(image_paths))
dataset = dataset.batch(2)
# 使用數據集
for images, bboxes, class_ids in dataset:
print(images.shape, bboxes, class_ids)
在上面的示例中,首先定義了圖片路徑和標注文件的列表。然后定義了一個parse_function函數來處理每個樣本的圖片和標注信息,最后使用tf.data.Dataset類的from_tensor_slices方法構建數據集,并利用map方法將parse_function函數應用到每個樣本上。接著通過shuffle和batch方法對數據集進行亂序和分批處理,最后可以使用數據集迭代獲取圖片,邊框和類別信息。