经典卷积模型回顾29—YOLOV2实现垃圾检测（Tensorflow2.0）

1. 下载yolov2模型的权重和配置文件

```

!wget https://github.com/pjreddie/darknet/raw/master/cfg/yolov2.cfg

!wget https://pjreddie.com/media/files/yolov2.weights

```

2. 安装相关Python库

```

!pip install tensorflow==2.0.0 matplotlib opencv-python

```

3. 加载模型权重和配置文件，生成模型

```

import cv2

import numpy as np

import tensorflow as tf

 

# 定义相关参数

model_cfg = 'yolov2.cfg'

model_weights = 'yolov2.weights'

image_size = 416

classes = ['bottle', 'can', 'box']

 

# 读取配置文件和权重文件

with open(model_cfg, 'r') as file:

    model_str = file.read()

model_data = tf.keras.utils.get_file(fname=model_weights, origin='https://pjreddie.com/media/files/yolov2.weights')

yolo_model = tf.keras.models.model_from_config(model_str)

 

# 读取权重文件中的参数，并设置到模型中对应的层

# 注意：yolov2.weights是Darknet的权重格式，需要进行转换

def load_weights(model, weights_file):

    with open(weights_file, 'rb') as f:

        major, minor, revision, seen, _ = np.fromfile(f, dtype=np.int32, count=5)

 

        for layer in model.layers:

            if any([s in layer.name for s in ['conv2d', 'batch_normalization', 'leaky']]):

                print("loading layer:", layer.name)

                weights_shape = [w.shape for w in layer.get_weights()]

                if len(weights_shape) == 4:

                    # darknet shape (out_dim, in_dim, height, width)

                    # tensorflow shape (height, width, in_dim, out_dim)

                    weights_shape = [weights_shape[3], weights_shape[2], weights_shape[0], weights_shape[1]]

                elif len(weights_shape) == 2:

                    weights_shape = list(reversed(weights_shape))

                else:

                    raise ValueError('invalid weights')

                weights_size = np.product(weights_shape)

 

                # read weights

                weights = np.fromfile(f, dtype=np.float32, count=weights_size)

                weights = weights.reshape(weights_shape)

 

                # copy weights to model

                layer.set_weights([weights])

fcn_model = yolo_model.get_layer('model')

load_weights(fcn_model, model_data)

 

# 定义输出层（yolo_detection），并连接到模型中

input_image = yolo_model.input

outputs = []

for i in range(3):

    output = yolo_model.get_layer('leaky_re_lu_' + str(17 + 4 * i)).output

    output = yolo_model.get_layer('conv2d_' + str(18 + 4 * i)).output(output)

    outputs.append(output)

yolo_detection = tf.keras.layers.concatenate(outputs, axis=1, name='yolo_detection')

yolo_model = tf.keras.models.Model(input_image, yolo_detection)

```

 

4. 定义辅助函数，用于处理检测结果

 

```

def sigmoid(x):

    return 1 / (1 + np.exp(-x))

 

def softmax(x):

    return np.exp(x) / np.sum(np.exp(x))

 

def non_max_suppression(boxes, confidences, threshold):

    indices = np.argsort(-confidences)

    selected_indices = []

    while len(indices) > 0:

        i = indices[0]

        selected_indices.append(i)

        if len(indices) == 1:

            break

        ious = []

        for j in indices[1:]:

            iou = box_iou(boxes[i], boxes[j])

            ious.append(iou)

        ious = np.array(ious)

        filtered_indices = np.where(ious < threshold)[0] + 1

        indices = indices[filtered_indices]

    return selected_indices

 

def box_iou(box1, box2):

    x1, y1, w1, h1 = box1

    x2, y2, w2, h2 = box2

    intersection = min(x1 + w1, x2 + w2) - max(x1, x2)

    intersection *= min(y1 + h1, y2 + h2) - max(y1, y2)

    area1, area2 = w1 * h1, w2 * h2

    iou = intersection / (area1 + area2 - intersection)

    return iou

```

 

5. 定义检测函数，并加载测试图片进行检测

 

```

def detect_image(image_file):

    # 读取输入图片，并调整大小

    image = cv2.imread(image_file)

    height, width, _ = image.shape

    new_height = height * min(image_size / height, image_size / width)

    new_width = width * min(image_size / height, image_size / width)

    resized_image = cv2.resize(image, (int(new_width), int(new_height)), cv2.INTER_CUBIC)

    canvas = np.full((image_size, image_size, 3), 128)

    canvas[(image_size - int(new_height)) // 2:(image_size - int(new_height)) // 2 + int(new_height),

           (image_size - int(new_width)) // 2:(image_size - int(new_width)) // 2 + int(new_width)] = resized_image

    # 将图片转换成TensorFlow的输入格式

    input_image = canvas[..., ::-1].astype(np.float32)

    input_image /= 255.

    input_image = np.expand_dims(input_image, 0)

    # 前向推理

    outputs = yolo_model.predict(input_image)

    # 解析检测结果

    result_boxes = []

    result_classes = []

    result_scores = []

    for i in range(len(outputs)):

        output = outputs[i][0]

        grid_size = output.shape[0]

        stride = int(image_size / grid_size)

        anchors = [[10, 14], [23, 27], [37, 58], [81, 82], [135, 169], [344, 319]]

        boxes = []

        box_scores = []

        for row in range(grid_size):

            for col in range(grid_size):

                channel = 5 + len(classes)

                scores = output[row, col, :, :channel]

                scores = softmax(scores)

                class_id = np.argmax(scores[:, 5:], axis=-1)

                class_id = class_id.astype(np.float32)

                obj_conf = sigmoid(scores[:, 4])

                scores = obj_conf * scores[:, 5:]

                scores *= scores > 0.3

                for i in range(scores.shape[0]):

                    if np.max(scores[i, :]) < 1e-6:

                        continue

                    x, y, w, h = output[row, col, i, :4]

                    x = (col + sigmoid(x)) * stride

                    y = (row + sigmoid(y)) * stride

                    w = anchors[2 * i + 1][0] * np.exp(w) * stride

                    h = anchors[2 * i + 1][1] * np.exp(h) * stride

                    box = [x - w / 2, y - h / 2, w, h]

                    boxes.append(box)

                    box_scores.append(scores[i, :])

        boxes = np.array(boxes)

        box_scores = np.array(box_scores)

        indices = non_max_suppression(boxes, box_scores[:, 0], 0.4)

        selected_boxes = boxes[indices]

        selected_classes = np.argmax(box_scores[indices, 1:], axis=-1)

        selected_classes = np.array([classes[i] for i in selected_classes])

        selected_scores = np.max(box_scores[indices, 1:], axis=-1)

        result_boxes.append(selected_boxes)

        result_classes.append(selected_classes)

        result_scores.append(selected_scores)

    result_boxes = np.concatenate(result_boxes, axis=0)

    result_classes = np.concatenate(result_classes, axis=0)

    result_scores = np.concatenate(result_scores, axis=0)

    # 在原图上绘制检测结果

    for i in range(len(result_boxes)):

        x, y, w, h = result_boxes[i]

        class_id = result_classes[i]

        label = '%s %.2f' % (class_id, result_scores[i])

        color = [0, 255, 0]

        cv2.rectangle(image, (int(x), int(y)), (int(x + w), int(y + h)), color, 2)

        cv2.putText(image, label, (int(x), int(y) - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)

    cv2.imwrite('result.jpg', image)

# 加载测试图片进行检测

detect_image('test.jpg')

```

运行以上代码后，将会在当前目录下生成一张result.jpg图片，该图片上会将检测结果标注出来。