TensorFlow下使用YOLOv1训练+测试自己的数据集

TensorFlow下使用YOLOv1训练+测试自己的数据集

 环境：

win10 + cuda8.0 + cudnn v5.1 + python3.5 + tensorflow-gpu 1.2.1

1、获取数据集

创建..\yolo_tensorflow-master\data\pascal_voc\路径，复制VOCdevkit文件放在pascal_voc目录下。

2、修改文件

① ..\yolo_tensorflow-master\yolo\config.py文件 

# 第20行
CLASSES = ['ali', 'jg3', 'ddg1000', 'jilong'] # 修改为自己的类别

 ② ..\yolo_tensorflow-master\utils\pascal_voc.py文件

# 第32行
        labels = np.zeros(
            (self.batch_size, self.cell_size, self.cell_size, 9)) # 25修改为5+类别数


# 第126行
        label = np.zeros((self.cell_size, self.cell_size, 9)) # 修改为5+类别数

③ ..\yolo_tensorflow-master\train.py文件

# 第80行
                    log_str = '{} Epoch: {}, Step: {}, Learning rate: {}, Loss: {:5.3f}\nSpeed: {:.3f}s/iter, Load: {:.3f}s/iter, Remain: {}'.format(
                        datetime.datetime.now().strftime('%m-%d %H:%M:%S'),

3、开始训练，运行 train.py文件。

4、参考

【YOLO初探】yolo_tensorflow_v1训练自己的数据集 

5、测试

① 将..\yolo_tensorflow-master\data\pascal_voc\output\2019_03_28_15_56\目录下（2019_03_28_15_56文件夹在模型开始训练时自动生成）保存的最后的模型文件（下图1）复制到..\yolo_tensorflow-master\data\weights\目录下（下图2），并在名称中间添加“.ckpt”重新命名。

② 修改..\yolo_tensorflow-master\test.py文件。

将待测试图片放入..\yolo_tensorflow-master\test\目录下。

# 将189行的默认权重文件改为自己训练好的模型
def main():
    parser = argparse.ArgumentParser()
    # parser.add_argument('--weights', default="YOLO_small.ckpt", type=str)
    parser.add_argument('--weights', default="yolo-15000.ckpt", type=str)


# 将207行需要测试的图片名称修改为自己的图片名称
    # detect from image file
    imname = 'test/000004.jpg'
    detector.image_detector(imname)

③ 运行test.py文件进行测试。