/keypoint-resnet50

Based on ResNet50, detect keypoint, and use keypoint perspective transform image.

Primary LanguagePythonMIT LicenseMIT

基于ResNet的关键点检测

项目简介

  • 本项目实现对单一目标的关键点检测
  • 主要用于检测矩形物体的四个角点，比如车牌、名片等
  • 检测到四个角点的关键点后，可使用透视变换将图形矫正

环境构建

  1. 创建虚拟环境
conda create -n keypoint python=3.7
conda activate keypoint
  1. 安装cuda和cudnn
conda install cudatoolkit=10.2
conda install cudnn
  1. 安装pytorch
pip install torch==1.7.0
pip install torchvision==0.8.0
  1. 安装项目依赖的其他包
pip install -r requirements.txt

数据整理

  1. 标注
  • 使用labelme对目标的四个角点做关键点标注，得到json格式的标注文件
  1. 数据集整理
  • 按照如下格式存放数据，项目中存放了若干示例的图像和标签供参考
项目目录/data
└── car_plate  # 数据集名称
    ├── train
    |   ├── images  # 训练集图片
    |   └── labels  # 训练集标签，json格式标注文件
    └── val
        ├── images  # 验证集图片
        └── labels  # 验证集标签，json格式标注文件

训练

python train.py --batch-size 4 --epochs 500
  • 出现如下日志，则训练正常
INFO: Namespace(batch_size=4, data_dir='./data/car_plate', epochs=500, heatmap_size=112, image_size=448, kpt_num=4, lr=0.001, model_save_dir='./model/', save_interval=50, sigma=3.0, weight_decay=0.0008)
Number of images in train: 2148
Number of images in val: 544
INFO: Using device cuda:0
INFO: Start training.
INFO: Epoch 1 start.
Train loss: 0.002156: 100%|███████████████████████████████████████████████| 537/537 [00:04<00:00,  5.32it/s]
INFO: Epoch: 1 train done, Loss: 0.002156, Accuracy: 19.9811%.
Test loss: 0.002241: 100%|████████████████████████████████████████████████| 136/136 [00:00<00:00,  8.16it/s]
INFO: Epoch: 1 valid done, Loss: 0.002241, Accuracy: 12.5000%.
  • 可以通过 tensorboard 查看 loss 变化
tensorboard --logdir model/

推理

  • 启动推理
python inference.py
  • 效果展示
关键点检测结果 图形矫正结果
keypoint perspective