Tiny_Yolov1_VOC2007目标检测

Tiny_Yolov1_VOC2007目标检测

相较于YOLO模型,Tiny_YOLO版本将网络压缩了许多,不管是训练还是移植速度都比较快,更加适用于业界应用。

一、 Tiny_Yolov1结构

网络输入:(448,448,3)

第一轮处理:
(1)Conv2D(16, (3, 3), padding=‘same’, use_bias=False)
(2)BatchNormalization
(3)LeakyReLU(alpha=0.1)
(4)MaxPooling2D((2, 2), strides=(2, 2), padding=‘same’)

第二轮处理:
(1)Conv2D(32, (3, 3), padding=‘same’, use_bias=False)
(2)BatchNormalization
(3)LeakyReLU(alpha=0.1)
(4)MaxPooling2D((2, 2), strides=(2, 2), padding=‘same’)

第三轮处理:
(1)Conv2D(64, (3, 3), padding=‘same’, use_bias=False)
(2)BatchNormalization
(3)LeakyReLU(alpha=0.1)
(4)MaxPooling2D((2, 2), strides=(2, 2), padding=‘same’)

第四轮处理:
(1)Conv2D(128, (3, 3), padding=‘same’, use_bias=False)
(2)BatchNormalization
(3)LeakyReLU(alpha=0.1)
(4)MaxPooling2D((2, 2), strides=(2, 2), padding=‘same’)

第五轮处理:
(1)Conv2D(256, (3, 3), padding=‘same’, use_bias=False)
(2)BatchNormalization
(3)LeakyReLU(alpha=0.1)
(4)MaxPooling2D((2, 2), strides=(2, 2), padding=‘same’)

第六轮处理:
(1)Conv2D(512, (3, 3), padding=‘same’, use_bias=False)
(2)BatchNormalization
(3)LeakyReLU(alpha=0.1)
(4)MaxPooling2D((2, 2), strides=(2, 2), padding=‘same’)

第七轮处理:
(1)Conv2D(1024, (3, 3), padding=‘same’, use_bias=False)
(2)BatchNormalization
(3)LeakyReLU(alpha=0.1)

第八轮处理:
(1)Conv2D(256, (3, 3), padding=‘same’, use_bias=False)
(2)BatchNormalization
(3)LeakyReLU(alpha=0.1)

第九轮处理:
(1)Flatten()
(2)Dropout(0.5)
(3)Dense(1470, activation=‘relu’)
(4)Dropout(0.5)
(5)Dense(1470, activation=‘linear’)

第十轮处理:
(1)my_reshape((7, 7, 30))

二、VOC2007数据集上的训练过程

第一轮训练过程:epoch = 5,optimizer=‘adam’, train loss从478.6下降到55.1,val loss从88.3下降到75.1,此时val loss训练已经达到瓶颈,无法再继续下降。

此时要想继续训练使得val loss下降已经颇有困难,得非常具有技巧性。
在这里插入图片描述
第二轮训练过程:epoch = 10,sgd = optimizers.SGD(lr=1e-4, momentum=0.9),train loss从52.4下降到47.3,val loss从74.1下降到73.1。

val loss下降非常缓慢,也非常困难。

在这里插入图片描述
第三轮训练过程:epoch=30,sgd = optimizers.SGD(lr=1e-5, momentum=0.9),train loss从48.0下降到46.4,val loss从73.3下降到72.5。

此轮训练我又进一步调低了学习率,只期望val loss能下降就好。

在这里插入图片描述
第四轮训练过程:epoch=500,sgd = optimizers.SGD(lr=1e-5, momentum=0.9),train loss从47.2下降到41.0,val loss从72.7下降到71.7。

val loss下降越来越困难了,似乎已经达到了极限,我没办法让训练效果再好了。

在这里插入图片描述

三、模型检测效果

我从测试集中随意选取10张图片,检测效果如下:

在这里插入图片描述
在这里插入图片描述
在这里插入图片描述
在这里插入图片描述
在这里插入图片描述
在这里插入图片描述
在这里插入图片描述
在这里插入图片描述
在这里插入图片描述

在这里插入图片描述

这个项目前前后后花费了我快一个月时间,我已经尽力了,YOLOv1模型检测我只能做到这种效果了,以后有想法再继续实验探究。

四、深入思考

Ques1:自己发现之前项目思路中的一个bug。

之前自己写过一个YOLOv1_VOC2007检测模型,网络训练了一个礼拜的时间,最后在训练集上检测效果非常好,而在测试集上的检测效果一塌糊涂。

本来我以为是过拟合问题,训练集loss很低,但验证集loss却降不下去。在尝试添加正则项、简化模型结构、dropout都不起作用后,暂时放弃了这个问题的处理。这段时间我才突然反应过来,根本不是过拟合问题,甚至可以说是欠拟合问题。

我在网络搭建的源代码中,Conv2D、BN层设置了trainable = False,但网络训练却是从头开始训练的。也就是说,卷积层、BN层的权重参数是随机生成的,并且一直无法被更新。我训练了一个礼拜的网络,其实只是在训练最后一层2亿个参数的全卷积。卷积、池化层并没有提取到正确的图片特征,而全连接层的训练只是在不断强行拟合,自然测试集上的效果会一塌糊涂。

实际上,要么我把所有的trainable设置为True,从头开始自己训练模型;要么从网上download已经训练好的部分网络层权重,再把trainable设置为 False。

Ques2:迁移学习。

迁移学习不是一种算法而是一种机器学习思想,应用到深度学习就是微调(Fine-tune)。修改网络模型结构,选择性载入网络模型权重,再用自己的数据集重新训练模型。微调能够快速训练好一个模型,用相对较小的数据量达到不错的结果。

网络训练有两种策略:一种是白手起家从头搭建模型进行训练,一种是通过预训练模型进行训练。

为什么要迁移学习?
(1)站在巨人的肩膀上。前人花很大精力训练出来的模型,大概率会比自己从零训练的效果更好。
(2)训练成本可以很低。迁移学习后期的训练成本非常低,用CPU都完全无压力。
(3)适用于小数据集。对于数据集本身很小(几千张图片)的情况,从头开始训练具有几千万参数的大型神经网络是不现实的,我们可以借用上大型神经网络的超强特征提取能力。

迁移学习有以下3种处理方式:

(1)Transfer Learning:冻结预训练模型的全部卷积层,只训练自己特定的全连接层。
(2)Extract Feature Vector:先计算出预训练模型的卷积层对所有训练、测试数据的特征向量,并存储到硬盘中,然后抛开预训练模型,只训练自己设计的简配版全连接网络。
(3)Fine-tune:冻结预训练模型的部分卷积层(通常是靠近输入的多数卷积层),训练剩下的卷积层(通常是靠近输出的部分卷积层)和全连接层。

第三种是对前两种的补充,但并不一定能真的对模型有所提升,不要期待能有什么质的飞跃。

Ques3:YOLOv1网络训练技巧。

不知道是因为YOLO模型结构太过复杂,还是YOLO损失函数太过复杂,YOLOv1模型真是太难训练了,比前面的图片分类模型、语义分割模型都要复杂的多。

训练技巧:先设置一个sgd学习率训练一段时间,一旦出现train loss不断下降但val loss基本不变的情况,立马停止训练保留此时的权重,此时已经开始向错误方向进行权重优化了,如果再不停止后面训练会更加艰难。尝试将sgd学习率下降10倍,开始下一轮的网络训练。

五、源码

读取数据集:

import os
import xml.etree.ElementTree as ET


class_dictionary = {'aeroplane': 0, 'bicycle': 1, 'bird': 2, 'boat': 3, 'bottle': 4, 'bus': 5,
                    'car': 6, 'cat': 7, 'chair': 8, 'cow': 9, 'diningtable': 10, 'dog': 11,
                    'horse': 12, 'motorbike': 13, 'person': 14, 'pottedplant': 15, 'sheep': 16,
                    'sofa': 17, 'train': 18, 'tvmonitor': 19}
class_list = list(class_dictionary.keys())


def read_image_path():
    data_x = []
    filename = os.listdir('/home/zk/Desktop/yolov1/JPEGImages')
    filename.sort()
    for name in filename:
        path = '/home/zk/Desktop/yolov1/JPEGImages/' + name
        data_x.append(path)

    print('JPEGImages has been download ! ')
    return data_x


def read_coordinate_txt():
    data_y = []
    filename = os.listdir('/home/zk/Desktop/yolov1/Annotations')
    filename.sort()
    for name in filename:
        tree = ET.parse('/home/zk/Desktop/yolov1/Annotations/' + name)
        root = tree.getroot()

        coordinate = ''
        for obj in root.iter('object'):

            difficult = obj.find('difficult').text
            cls = obj.find('name').text
            if cls not in class_list or int(difficult) == 1:
                continue

            cls_id = class_list.index(cls)
            xml_box = obj.find('bndbox')
            x_min = int(xml_box.find('xmin').text)
            y_min = int(xml_box.find('ymin').text)
            x_max = int(xml_box.find('xmax').text)
            y_max = int(xml_box.find('ymax').text)

            loc = (str(x_min) + ',' + str(y_min) + ',' + str(x_max) + ',' + str(y_max) + ',' + str(cls_id) + '  ')
            coordinate = coordinate + loc

        data_y.append(coordinate)

    print('Object Coordinate has been download ! ')
    return data_y


def make_data():
    data_x = read_image_path()
    data_y = read_coordinate_txt()

    n = len(data_x)
    train_x = data_x[0:8000]
    train_y = data_y[0:8000]
    val_x = data_x[8000:9000]
    val_y = data_y[8000:9000]
    test_x = data_x[9000:n]
    test_y = data_y[9000:n]

    return train_x, train_y, val_x, val_y, test_x, test_y



损失函数:

import keras.backend as k


def iou(pre_min, pre_max, true_min, true_max):

    intersect_min = k.maximum(pre_min, true_min)    # batch * 7 * 7 * 2 * 2
    intersect_max = k.minimum(pre_max, true_max)    # batch * 7 * 7 * 2 * 2

    intersect_wh = k.maximum(intersect_max - intersect_min, 0.)    # batch * 7 * 7 * 2 * 2
    intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]    # batch * 7 * 7 * 2

    pre_wh = pre_max - pre_min    # batch * 7 * 7 * 2 * 2
    true_wh = true_max - true_min    # batch * 7 * 7 * 1 * 2
    pre_area = pre_wh[..., 0] * pre_wh[..., 1]    # batch * 7 * 7 * 2
    true_area = true_wh[..., 0] * true_wh[..., 1]   # batch * 7 * 7 * 1

    union_area = pre_area + true_area - intersect_area    # batch * 7 * 7 * 2
    iou_score = intersect_area / union_area    # batch * 7 * 7 * 2

    return iou_score


def yolo_loss(y_true, y_pre):
    true_class = y_true[..., :20]    # batch * 7 * 7 * 20
    true_confidence = y_true[..., 20]    # batch * 7 * 7
    true_confidence1 = k.expand_dims(true_confidence)    # batch * 7 * 7 * 1
    true_location = y_true[..., 21:25]    # batch * 7 * 7 * 4

    pre_class = y_pre[..., :20]    # batch * 7 * 7 * 20
    pre_confidence = y_pre[..., 20:22]    # batch * 7 * 7 * 2
    pre_location = y_pre[..., 22:30]    # batch * 7 * 7 * 8

    true_location1 = k.reshape(true_location, [-1, 7, 7, 1, 4])    # batch * 7 * 7 * 1 * 4
    pre_location1 = k.reshape(pre_location, [-1, 7, 7, 2, 4])    # batch * 7 * 7 * 2 * 4

    true_xy = true_location1[..., :2] * 448 / 7    # batch * 7 * 7 * 1 * 2
    true_wh = true_location1[..., 2:4] * 448     # batch * 7 * 7 * 1 * 2
    true_xy_min = true_xy - true_wh / 2    # batch * 7 * 7 * 1 * 2
    true_xy_max = true_xy + true_wh / 2    # batch * 7 * 7 * 1 * 2

    pre_xy = pre_location1[..., :2] * 448 / 7    # batch * 7 * 7 * 2 * 2
    pre_wh = pre_location1[..., 2:4] * 448    # batch * 7 * 7 * 2 * 2
    pre_xy_min = pre_xy - pre_wh / 2    # batch * 7 * 7 * 2 * 2
    pre_xy_max = pre_xy + pre_wh / 2  # batch * 7 * 7 * 2 * 2

    iou_score = iou(pre_xy_min, pre_xy_max, true_xy_min, true_xy_max)  # batch * 7 * 7 * 2
    best_score = k.max(iou_score, axis=3, keepdims=True)   # batch * 7 * 7 * 1
    box_mask = k.cast(iou_score >= best_score, k.dtype(iou_score))  # batch * 7 * 7 * 2

    no_object_loss = 0.5 * (1 - box_mask * true_confidence1) * k.square(0 - pre_confidence)
    object_loss = box_mask * true_confidence1 * k.square(1 - pre_confidence)
    confidence_loss = no_object_loss + object_loss
    confidence_loss = k.sum(confidence_loss)

    class_loss = true_confidence1 * k.square(true_class - pre_class)
    class_loss = k.sum(class_loss)

    box_mask1 = k.expand_dims(box_mask)    # batch * 7 * 7 * 2 * 1
    true_confidence2 = k.expand_dims(true_confidence1)    # batch * 7 * 7 * 1 * 1

    location_loss_xy = 5 * box_mask1 * true_confidence2 * k.square((true_xy - pre_xy) / 448)
    location_loss_wh = 5 * box_mask1 * true_confidence2 * k.square((k.sqrt(true_wh) - k.sqrt(pre_wh)) / 448)
    location_loss = k.sum(location_loss_xy) + k.sum(location_loss_wh)

    total_loss = confidence_loss + class_loss + location_loss

    return total_loss

网络结构:

import numpy as np
from keras.models import Sequential, Model, Input
from keras.layers import Dense, LeakyReLU, Flatten, Dropout
from keras.layers import Conv2D, MaxPooling2D
from keras.layers import BatchNormalization
from keras import layers
from keras.engine.topology import Layer
import keras.backend as k
from keras.regularizers import l2


class my_reshape(Layer):
    def __init__(self, target_shape, **kwargs):
        super(my_reshape, self).__init__(**kwargs)
        self.target_shape = target_shape    # (7, 7, 30)

    def compute_output_shape(self, input_shape):
        return (input_shape[0], ) + self.target_shape    # (batch, 7, 7, 30)

    def call(self, inputs, **kwargs):
        s = [self.target_shape[0], self.target_shape[1]]    # [7, 7]
        c = 20
        b = 2

        idx1 = s[0] * s[1] * c    # 7 * 7 * 20 = 980
        idx2 = idx1 + s[0] * s[1] * b    # 980 + 7 * 7 * 2 = 1078

        class_tensor = k.reshape(inputs[:, :idx1], (k.shape(inputs)[0],) + tuple([s[0], s[1], c]))
        class_tensor = k.softmax(class_tensor)    # shape = (batch, 7, 7, 20)

        confidence_tensor = k.reshape(inputs[:, idx1:idx2], (k.shape(inputs)[0],) + tuple([s[0], s[1], b]))
        confidence_tensor = k.sigmoid(confidence_tensor)    # shape = (batch, 7, 7, 2)

        location_tensor = k.reshape(inputs[:, idx2:], (k.shape(inputs)[0],) + tuple([s[0], s[1], b * 4]))
        location_tensor = k.sigmoid(location_tensor)    # shape = (batch, 7, 7, 8)

        outputs = k.concatenate([class_tensor, confidence_tensor, location_tensor])    # shape = (batch, 7, 7, 30)

        return outputs


def create_network():

    inputs = Input((448, 448, 3))

    x = Conv2D(16, (3, 3), padding='same', name='convolutional_0', use_bias=False,
               kernel_regularizer=l2(5e-4))(inputs)
    x = BatchNormalization(name='bnconvolutional_0')(x)
    x = LeakyReLU(alpha=0.1)(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), padding='same')(x)

    x = Conv2D(32, (3, 3), padding='same', name='convolutional_1', use_bias=False,
               kernel_regularizer=l2(5e-4))(x)
    x = BatchNormalization(name='bnconvolutional_1')(x)
    x = LeakyReLU(alpha=0.1)(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), padding='same')(x)

    x = Conv2D(64, (3, 3), padding='same', name='convolutional_2', use_bias=False,
               kernel_regularizer=l2(5e-4))(x)
    x = BatchNormalization(name='bnconvolutional_2')(x)
    x = LeakyReLU(alpha=0.1)(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), padding='same')(x)

    x = Conv2D(128, (3, 3), padding='same', name='convolutional_3', use_bias=False,
               kernel_regularizer=l2(5e-4))(x)
    x = BatchNormalization(name='bnconvolutional_3')(x)
    x = LeakyReLU(alpha=0.1)(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), padding='same')(x)

    x = Conv2D(256, (3, 3), padding='same', name='convolutional_4', use_bias=False,
               kernel_regularizer=l2(5e-4))(x)
    x = BatchNormalization(name='bnconvolutional_4')(x)
    x = LeakyReLU(alpha=0.1)(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), padding='same')(x)

    x = Conv2D(512, (3, 3), padding='same', name='convolutional_5', use_bias=False,
               kernel_regularizer=l2(5e-4))(x)
    x = BatchNormalization(name='bnconvolutional_5')(x)
    x = LeakyReLU(alpha=0.1)(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), padding='same')(x)

    x = Conv2D(1024, (3, 3), padding='same', name='convolutional_6', use_bias=False,
               kernel_regularizer=l2(5e-4))(x)
    x = BatchNormalization(name='bnconvolutional_6')(x)
    x = LeakyReLU(alpha=0.1)(x)

    x = Conv2D(256, (3, 3), padding='same', name='convolutional_7', use_bias=False,
               kernel_regularizer=l2(5e-3))(x)
    x = BatchNormalization(name='bnconvolutional_7')(x)
    x = LeakyReLU(alpha=0.1)(x)

    x = Flatten()(x)
    x = Dropout(0.5)(x)

    x = Dense(1470, activation='relu', name='connected_0')(x)
    x = Dropout(0.5)(x)

    x = Dense(1470, activation='linear', name='connected_1')(x)
    outputs = my_reshape((7, 7, 30))(x)
    model = Model(inputs=inputs, outputs=outputs)

    return model


训练过程:

import numpy as np
from keras.models import load_model
import cv2
import tiny_yolov1_model
from yolo_loss import yolo_loss
from keras.utils import Sequence
import math
from keras.callbacks import ModelCheckpoint
from keras import optimizers


class SequenceData(Sequence):

    def __init__(self, data_x, data_y, batch_size):
        self.batch_size = batch_size
        self.data_x = data_x
        self.data_y = data_y
        self.indexes = np.arange(len(self.data_x))

    def __len__(self):
        return math.floor(len(self.data_x) / float(self.batch_size))

    def on_epoch_end(self):
        np.random.shuffle(self.indexes)

    def __getitem__(self, idx):

        batch_index = self.indexes[idx * self.batch_size:(idx + 1) * self.batch_size]
        batch_x = [self.data_x[k] for k in batch_index]
        batch_y = [self.data_y[k] for k in batch_index]

        x = np.zeros((len(batch_x), 448, 448, 3))
        y = np.zeros((len(batch_y), 7, 7, 25))

        for i in range(self.batch_size):

            img = cv2.imread(batch_x[i])
            obj_all = batch_y[i].strip().split()
            size = img.shape

            img1 = img / 255
            resize_img = cv2.resize(img1, (448, 448), interpolation=cv2.INTER_AREA)
            x[i, :, :, :] = resize_img

            for j in range(len(obj_all)):
                obj = obj_all[j].split(',')
                x1, y1, x2, y2 = [int(obj[0]), int(obj[1]), int(obj[2]), int(obj[3])]
                category = int(obj[4])

                center_x = (x1 + x2) / 2
                center_y = (y1 + y2) / 2
                w = x2 - x1
                h = y2 - y1

                grid_x = int(7 * center_x / size[1])
                grid_y = int(7 * center_y / size[0])

                center_x_ratio = center_x * 7 / size[1] - grid_x
                center_y_ratio = center_y * 7 / size[0] - grid_y
                w_ratio = w / size[1]
                h_ratio = h / size[0]

                y[i, grid_y, grid_x, category] = 1
                y[i, grid_y, grid_x, 20] = 1
                y[i, grid_y, grid_x, 21:25] = np.array([center_x_ratio, center_y_ratio, w_ratio, h_ratio])

        return x, y


# create model and train and save
def train_network(train_generator, validation_generator, epoch):

    model = tiny_yolov1_model.create_network()
    model.load_weights('/home/zk/Desktop/yolov1/raw_tiny-yolov1.hdf5', by_name=True)

    for i in range(31):
        model.layers[i].trainable = False
        print(model.layers[i])

    model.summary()
    model.compile(loss=yolo_loss, optimizer='adam')

    checkpoint = ModelCheckpoint('/home/zk/Desktop/yolov1/best_weights.hdf5', monitor='val_loss',
                                 save_weights_only=True, save_best_only=True)

    model.fit_generator(
        train_generator,
        steps_per_epoch=len(train_generator),
        epochs=epoch,
        validation_data=validation_generator,
        validation_steps=len(validation_generator),
        callbacks=[checkpoint]
    )

    model.save_weights('first_weights.hdf5')


# Load the partially trained model and continue training and save
def load_network_then_train(train_generator, validation_generator, epoch, input_name, output_name):

    model = tiny_yolov1_model.create_network()
    model.load_weights(input_name)

    for i in range(31):
        model.layers[i].trainable = False
        print(model.layers[i])

    model.summary()

    sgd = optimizers.SGD(lr=1e-5, momentum=0.9)
    model.compile(loss=yolo_loss, optimizer=sgd)
    checkpoint = ModelCheckpoint('/home/zk/Desktop/yolov1/best_weights.hdf5', monitor='val_loss',
                                 save_weights_only=True, save_best_only=True)

    model.fit_generator(
        train_generator,
        steps_per_epoch=len(train_generator),
        epochs=epoch,
        validation_data=validation_generator,
        validation_steps=len(validation_generator),
        callbacks=[checkpoint]
    )

    model.save_weights(output_name)

主函数调用:

import read_data_path as rp
import numpy as np
import cv2
import train as tr
from train import SequenceData
import yolo_loss
import tiny_yolov1_model

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0"


class_dictionary = {'aeroplane': 0, 'bicycle': 1, 'bird': 2, 'boat': 3, 'bottle': 4, 'bus': 5,
                    'car': 6, 'cat': 7, 'chair': 8, 'cow': 9, 'diningtable': 10, 'dog': 11,
                    'horse': 12, 'motorbike': 13, 'person': 14, 'pottedplant': 15, 'sheep': 16,
                    'sofa': 17, 'train': 18, 'tvmonitor': 19}
class_list = list(class_dictionary.keys())

if __name__ == "__main__":

    train_x, train_y, val_x, val_y, test_x, test_y = rp.make_data()
    train_generator = SequenceData(train_x, train_y, 32)
    validation_generator = SequenceData(val_x, val_y, 32)

    # tr.train_network(train_generator, validation_generator, epoch=5)

    # tr.load_network_then_train(train_generator, validation_generator, epoch=10,
    #                            input_name='first_weights.hdf5', output_name='second_weights.hdf5')
    # tr.load_network_then_train(train_generator, validation_generator, epoch=30,
    #                            input_name='second_weights.hdf5', output_name='third_weights.hdf5')
    # tr.load_network_then_train(train_generator, validation_generator, epoch=500,
    #                            input_name='third_weights.hdf5', output_name='fourth_weights.hdf5')

    for i in range(len(test_x)-100, len(test_x)-90):

        img1 = cv2.imread(test_x[i])
        size = img1.shape

        img2 = img1 / 255
        img3 = cv2.resize(img2, (448, 448), interpolation=cv2.INTER_AREA)
        img4 = img3[np.newaxis, :, :, :]

        model = tiny_yolov1_model.create_network()
        model.compile(loss=yolo_loss, optimizer='adam')

        model.load_weights('fourth_weights.hdf5')
        pre = model.predict(img4)
        pre1 = pre[0]

        candidate = []
        for j in range(7):
            for k in range(7):
                pre1[j, k, 22] = pre1[j, k, 22] * 448 / 7 + k * 448 / 7
                pre1[j, k, 23] = pre1[j, k, 23] * 448 / 7 + j * 448 / 7
                pre1[j, k, 24] = pre1[j, k, 24] * 448
                pre1[j, k, 25] = pre1[j, k, 25] * 448
                pre1[j, k, 26] = pre1[j, k, 26] * 448 / 7 + k * 448 / 7
                pre1[j, k, 27] = pre1[j, k, 27] * 448 / 7 + j * 448 / 7
                pre1[j, k, 28] = pre1[j, k, 28] * 448
                pre1[j, k, 29] = pre1[j, k, 29] * 448

                if pre1[j, k, 20] > 0.5 or pre1[j, k, 21] > 0.5:

                    if pre1[j, k, 20] > pre1[j, k, 21]:
                        x1 = pre1[j, k, 22] - pre1[j, k, 24] / 2
                        y1 = pre1[j, k, 23] - pre1[j, k, 25] / 2
                        x2 = pre1[j, k, 22] + pre1[j, k, 24] / 2
                        y2 = pre1[j, k, 23] + pre1[j, k, 25] / 2
                        category = np.argmax(pre1[j, k, 0:20])
                        confidence = pre1[j, k, 20] * np.max(pre1[j, k, 0:20])

                        x1 = x1 / 448 * size[1]
                        y1 = y1 / 448 * size[0]
                        x2 = x2 / 448 * size[1]
                        y2 = y2 / 448 * size[0]

                        candidate.append([x1, y1, x2, y2, category, confidence])

                    else:
                        x1 = pre1[j, k, 26] - pre1[j, k, 28] / 2
                        y1 = pre1[j, k, 27] - pre1[j, k, 29] / 2
                        x2 = pre1[j, k, 26] + pre1[j, k, 28] / 2
                        y2 = pre1[j, k, 27] + pre1[j, k, 29] / 2
                        category = np.argmax(pre1[j, k, 0:20])
                        confidence = pre1[j, k, 21] * np.max(pre1[j, k, 0:20])

                        x1 = x1 / 448 * size[1]
                        y1 = y1 / 448 * size[0]
                        x2 = x2 / 448 * size[1]
                        y2 = y2 / 448 * size[0]

                        candidate.append([x1, y1, x2, y2, category, confidence])

        candidate = np.array(candidate)

        for num in range(len(candidate)):
            a1 = int(candidate[num, 0])
            b1 = int(candidate[num, 1])
            a2 = int(candidate[num, 2])
            b2 = int(candidate[num, 3])
            index = int(candidate[num, 4])
            pre_class = class_list[index]
            confidence = str(candidate[num, 5])

            cv2.rectangle(img1, (a1, b1), (a2, b2), (0, 0, 255), 2)
            cv2.putText(img1, pre_class, (a1, b1), 1, 1, (0, 0, 255))
            cv2.putText(img1, confidence, (a1, b2), 1, 1, (0, 0, 255))

        cv2.namedWindow("Image")
        cv2.imshow("Image", img1)
        cv2.waitKey(0)

        cv2.imwrite('/home/zk/Desktop/yolov1/demo/' + str(i) + '.jpg', img1)


六、相关链接

如果代码跑不通,或者想直接使用训练好的模型,可以去下载项目链接:
https://blog.csdn.net/Twilight737

  • 0
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值