我们用 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

img1.jpg

768

480

3

<object>

defect

100

150

200

250

</object>

· OK 样本（无缺陷）XML 文件示例：

xml

img2.jpg

768

480

3

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 构建并训练了一个目标检测模型，该模型能够处理包含缺陷和没有缺陷的图片。关键在于正确地处理没有目标的图片，并确保数据格式和类别索引的一致性。