写一个CNN的非极大抑制和训练

下面是一个简单的CNN训练和非极大抑制代码示例，使用了TensorFlow库：

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.layers import Lambda
import numpy as np
import cv2

# 构建CNN模型
model = Sequential()
model.add(Conv2D(32, (3,3), activation='relu', input_shape=(28,28,1)))
model.add(MaxPooling2D((2,2)))
model.add(Conv2D(64, (3,3), activation='relu'))
model.add(MaxPooling2D((2,2)))
model.add(Conv2D(64, (3,3), activation='relu'))
model.add(Flatten())
model.add(Dense(64, activation='relu'))
model.add(Dense(10, activation='softmax'))

# 编译模型
model.compile(optimizer=Adam(lr=0.001), loss='categorical_crossentropy', metrics=['accuracy'])

# 加载数据集
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

# 将图像转换为float类型并进行归一化
x_train = x_train.astype('float32') / 255
x_test = x_test.astype('float32') / 255

# 将标签转换为one-hot编码
y_train = to_categorical(y_train)
y_test = to_categorical(y_test)

# 对训练数据进行增强
datagen = ImageDataGenerator(rotation_range=10, zoom_range=0.1, width_shift_range=0.1, height_shift_range=0.1)
datagen.fit(x_train)

# 训练模型
model.fit(datagen.flow(x_train, y_train, batch_size=32), epochs=10, validation_data=(x_test, y_test))

# 非极大抑制
def non_max_suppression(boxes, scores, threshold):
    # 计算框的面积
    areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    # 根据得分从高到低排序
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        # 取得分最高的框
        i = order[0]
        keep.append(i)

        # 计算当前框与其它框的交集面积
        xx1 = np.maximum(boxes[i, 0], boxes[order[1:], 0])
        yy1 = np.maximum(boxes[i, 1], boxes[order[1:], 1])
        xx2 = np.minimum(boxes[i, 2], boxes[order[1:], 2])
        yy2 = np.minimum(boxes[i, 3], boxes[order[1:], 3])
        w = np.maximum(0, xx2 - xx1)
        h = np.maximum(0, yy2 - yy1)
        inter = w * h

        # 计算重叠度
        overlap = inter / (areas[i] + areas[order[1:]] - inter)

        # 保留重叠度小于阈值的框
        inds = np.where(overlap <= threshold)[0]
        order = order[inds + 1]

    return keep

# 测试模型
img = cv2.imread('test.jpg')
img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img_gray = cv2.resize(img_gray, (28, 28))
img_gray = img_gray.astype('float32') / 255
img_gray = np.expand_dims(img_gray, axis=2)
img_gray = np.expand_dims(img_gray, axis=0)
predictions = model.predict(img_gray)
class_idx = np.argmax(predictions)
class_output = model.output[:, class_idx]
last_conv_layer = model.get_layer('conv2d_3')
grads = K.gradients(class_output, last_conv_layer.output)[0]
pooled_grads = K.mean(grads, axis=(0, 1, 2))
iterate = K.function([model.input], [pooled_grads, last_conv_layer.output[0]])
pooled_grads_value, conv_layer_output_value = iterate([img_gray])
for i in range(64):
    conv_layer_output_value[:, :, i] *= pooled_grads_value[i]
heatmap = np.mean(conv_layer_output_value, axis=-1)
heatmap = np.maximum(heatmap, 0)
heatmap /= np.max(heatmap)
heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]))
heatmap = np.uint8(255 * heatmap)
heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
output_image = cv2.addWeighted(img, 0.5, heatmap, 0.5, 0)
cv2.imwrite('output.jpg', output_image)

# 非极大抑制
boxes = np.array([[40, 40, 100, 100], [50, 50, 120, 120], [60, 60, 130, 130]])
scores = np.array([0.9, 0.75, 0.8])
threshold = 0.5
keep = non_max_suppression(boxes, scores, threshold)
print(keep)

在上面的代码中，我们首先定义了一个简单的CNN模型，并使用Keras库进行训练。然后，我们使用OpenCV库的函数对一张测试图片进行预测，并生成热力图。接着，我们定义了一个非极大抑制的函数，用于过滤掉重叠度高的框。最后，我们对一组测试框进行非极大抑制，并输出保留的框的索引。