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文件进行测试。