我们用 TensorFlow 搭建和训练目标检测模型的代码，以便处理包含缺陷（NG样本）和正常（OK样本）的图片。我们将使用 TensorFlow 2 和 Keras API 来实现。
下面，我将为您提供一个完整的代码示例，包括数据准备、模型搭建和训练过程。这个示例将演示如何处理有缺陷和无缺陷的样本，并在训练过程中正确地处理没有目标的图片。
1. 环境准备
首先，确保您已经安装了以下软件和库：
· Python 3.6+
· TensorFlow 2.x
· Pillow（用于图像处理）
· KerasCV（Keras 计算机视觉组件，包含目标检测模型）
使用以下命令安装必要的库：
bash
pip install tensorflow
pip install pillow
pip install keras-cv --upgrade
2. 数据准备
2.1 数据集组织
假设您的数据集按以下结构组织：
dataset/
├── images/
│ ├── img1.jpg
│ ├── img2.jpg
│ ├── img3.jpg
│ └── ...
├── annotations/
│ ├── img1.xml
│ ├── img2.xml
│ ├── img3.xml
│ └── ...
· images/ 文件夹：存放所有的图片，包括 OK 和 NG 样本。
· annotations/ 文件夹：存放对应的标注文件，格式为 Pascal VOC 的 XML 文件。
o 对于 OK 样本（无缺陷），对应的 XML 文件中 没有 <object> 标签。
o 对于 NG 样本（有缺陷），XML 文件中包含缺陷的位置标注。
2.2 标注示例
· NG 样本（有缺陷）XML 文件示例：
xml
<annotation>
<filename>img1.jpg</filename>
<size>
<width>768</width>
<height>480</height>
<depth>3</depth>
</size>
<object>
<name>defect</name>
<bndbox>
<xmin>100</xmin>
<ymin>150</ymin>
<xmax>200</xmax>
<ymax>250</ymax>
</bndbox>
</object> </annotation>
· OK 样本（无缺陷）XML 文件示例：
xml
<annotation>
<filename>img2.jpg</filename>
<size>
<width>768</width>
<height>480</height>
<depth>3</depth>
</size>
<!-- 没有 <object> 标签 --> </annotation>
2.3 导入库
python
import os import xml.etree.ElementTree as ET import tensorflow as tf import numpy as np from PIL import Image import keras_cv from tensorflow import keras
2.4 定义参数和类别
python
# 定义图像尺寸
IMAGE_HEIGHT = 480
IMAGE_WIDTH = 768
# 定义类别列表（包括背景类）
CLASS_NAMES = ['bg', 'defect'] # 'bg' 表示背景类
NUM_CLASSES = len(CLASS_NAMES)
2.5 数据加载和解析
2.5.1 定义解析标注文件的函数
python
def parse_annotation(annotation_path):
tree = ET.parse(annotation_path)
root = tree.getroot()
filename = root.find('filename').text
size = root.find('size')
width = int(size.find('width').text)
height = int(size.find('height').text)
# 初始化边界框和类别列表
boxes = []
labels = []
# 查找所有的 object 标签
for obj in root.findall('object'):
name = obj.find('name').text
label = CLASS_NAMES.index(name)
bndbox = obj.find('bndbox')
xmin = int(bndbox.find('xmin').text) / width
ymin = int(bndbox.find('ymin').text) / height
xmax = int(bndbox.find('xmax').text) / width
ymax = int(bndbox.find('ymax').text) / height
boxes.append([ymin, xmin, ymax, xmax]) # 坐标格式为 y_min, x_min, y_max, x_max
labels.append(label)
return filename, boxes, labels
2.5.2 加载数据集
python
# 数据集路径
DATASET_PATH = 'dataset/' # 请根据实际情况调整路径
# 获取图像和标注文件列表
image_dir = os.path.join(DATASET_PATH, 'images')
annotation_dir = os.path.join(DATASET_PATH, 'annotations')
image_files = sorted(os.listdir(image_dir))
annotation_files = sorted(os.listdir(annotation_dir))
# 构建数据集列表
dataset = []
for annotation_file in annotation_files:
annotation_path = os.path.join(annotation_dir, annotation_file)
filename, boxes, labels = parse_annotation(annotation_path)
image_path = os.path.join(image_dir, filename)
data = {
'image_path': image_path,
'boxes': np.array(boxes, dtype=np.float32),
'labels': np.array(labels, dtype=np.int32)
}
dataset.append(data)
2.6 创建 TensorFlow Dataset
python
def load_data(data):
image = Image.open(data['image_path']).convert('RGB')
image = image.resize((IMAGE_WIDTH, IMAGE_HEIGHT))
image = np.array(image) / 255.0 # 归一化到 [0,1]
boxes = data['boxes']
labels = data['labels']
# 如果没有目标，boxes 和 labels 为空，需要特殊处理
if boxes.size == 0:
boxes = np.array([[0, 0, 0, 0]], dtype=np.float32)
labels = np.array([0], dtype=np.int32) # 背景类
return image, {'boxes': boxes, 'classes': labels}
# 构建 TensorFlow Dataset
images = []
targets = []
for data in dataset:
image, target = load_data(data)
images.append(image)
targets.append(target)
# 将列表转换为张量
images = np.array(images, dtype=np.float32)
boxes = [target['boxes'] for target in targets]
classes = [target['classes'] for target in targets]
2.7 创建数据生成器
由于每个样本的目标数量不同，我们需要创建一个生成器来提供数据，同时处理不同长度的边界框和类别。
python
def data_generator():
for img, bxs, cls in zip(images, boxes, classes):
yield img, {'boxes': bxs, 'classes': cls}
# 定义输出签名
output_signature = (
tf.TensorSpec(shape=(IMAGE_HEIGHT, IMAGE_WIDTH, 3), dtype=tf.float32),
{
'boxes': tf.TensorSpec(shape=(None, 4), dtype=tf.float32),
'classes': tf.TensorSpec(shape=(None,), dtype=tf.int32)
}
)
# 创建 Dataset
dataset = tf.data.Dataset.from_generator(
data_generator,
output_signature=output_signature
)
# 定义批大小和预取
BATCH_SIZE = 8
dataset = dataset.shuffle(1000).padded_batch(
BATCH_SIZE,
padding_values=(
tf.constant(0, dtype=tf.float32),
{
'boxes': tf.constant(0, dtype=tf.float32),
'classes': tf.constant(0, dtype=tf.int32)
}
),
drop_remainder=True
).prefetch(buffer_size=tf.data.AUTOTUNE)
3. 构建模型
3.1 使用 KerasCV 的 RetinaNet 模型
KerasCV 提供了预训练的目标检测模型，如 RetinaNet。我们可以使用预训练的 backbone，并根据我们的数据集进行微调。
python
from keras_cv.models import RetinaNet
model = RetinaNet(
classes=NUM_CLASSES, # 类别数量
bounding_box_format='yxyx', # 边界框格式
backbone='resnet50', # 预训练的主干网络
backbone_weights='imagenet', # 加载在 ImageNet 上预训练的权重
include_rescaling=False # 我们已经在数据加载时进行了归一化
)
3.2 编译模型
python
# 定义优化器和学习率
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4)
# 编译模型
model.compile(optimizer=optimizer)
4. 模型训练
4.1 开始训练
python
EPOCHS = 10
history = model.fit(
dataset,
epochs=EPOCHS,
verbose=1
)
5. 模型评估与预测
5.1 预测并可视化结果
python
import matplotlib.pyplot as plt
def visualize_detections(image, y_pred, threshold=0.5):
boxes = y_pred['boxes'][0].numpy()
scores = y_pred['confidence'][0].numpy()
classes = y_pred['classes'][0].numpy()
# 过滤低于阈值的预测
idxs = scores > threshold
boxes = boxes[idxs]
scores = scores[idxs]
classes = classes[idxs]
plt.figure(figsize=(10, 10))
plt.imshow(image)
ax = plt.gca()
for box, score, cls in zip(boxes, scores, classes):
y_min, x_min, y_max, x_max = box
# 将坐标转换为图像尺寸
x_min *= IMAGE_WIDTH
x_max *= IMAGE_WIDTH
y_min *= IMAGE_HEIGHT
y_max *= IMAGE_HEIGHT
width = x_max - x_min
height = y_max - y_min
rect = plt.Rectangle((x_min, y_min), width, height,
fill=False, color='red', linewidth=2)
ax.add_patch(rect)
text = f'{CLASS_NAMES[cls]}: {score:.2f}'
ax.text(x_min, y_min - 10, text, color='red', fontsize=12)
plt.axis('off')
plt.show()
# 从数据集中获取一张测试图片
test_image, _ = next(iter(dataset.unbatch().take(1)))
test_image_np = test_image.numpy()
# 扩展维度以匹配模型输入
test_image_input = np.expand_dims(test_image_np, axis=0)
# 进行预测
y_pred = model.predict(test_image_input)
# 可视化检测结果
visualize_detections(test_image_np, y_pred, threshold=0.5)
6. 注意事项·
处理没有目标的图片：在数据准备时，对于没有目标的图片，我们将 boxes 设置为 [[0, 0, 0, 0]]，classes 设置为 [0]，其中 0 表示背景类。·
类别索引：在 KerasCV 中，类别索引从 0 开始，其中 0 通常表示背景类。因此，要确保在标注解析时正确地映射类别索引。
边界框格式：我们使用 yxyx 格式（[y_min, x_min, y_max, x_max]），并且坐标值是相对于图像尺寸的归一化值（在 0 到 1 之间）。
批处理数据：由于每个样本的目标数量不同，我们使用 padded_batch 并指定 padding_values，以确保批数据形状一致。
训练超参数：可以根据实际情况调整批大小、学习率和训练轮数，以获得更好的模型性能。
·
7. 总结
通过以上步骤，我们成功地使用 TensorFlow 和 KerasCV 构建并训练了一个目标检测模型，该模型能够处理包含缺陷和没有缺陷的图片。关键在于正确地处理没有目标的图片，并确保数据格式和类别索引的一致性。