Crowd Counting P2PNet 复现

最近用p2pnet进行了复现，还对自己做的数据集进行了训练。在复现的过程中的一些注意事项在此做了记录。
代码地址:https://github.com/TencentYoutuResearch/CrowdCounting-P2PNet
论文地址：https://arxiv.org/abs/2107.12746
上海人群数据集地址：https://pan.quark.cn/s/4146a95ac9a6
提取码：TJaZ

P2PNET网络图

测试：
在用p2pnet 测试自己的图像数据时，需要先下载vgg的训练好的模型，因为p2pnet是由vgg提取出的特征图上进行预测，计数的。原程序中没有给vgg的模型需要自己下载。
vgg_.py

"""
import torch
import torch.nn as nn


__all__ = [
    'VGG', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn',
    'vgg19_bn', 'vgg19',
]


model_urls = {
    'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth',
    'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth',
    'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth',
    'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth',
    'vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth',
    'vgg13_bn': 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth',
    'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth',
    'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth',
}


model_paths = {
    'vgg16_bn': 'weights/vgg16_bn-6c64b313.pth',
    'vgg16': '/apdcephfs/private_changanwang/checkpoints/vgg16-397923af.pth',

}

这里我们可以看到在model_paths中给了vgg16_bn和vgg16两种模型，我们只需要下载其中一种模型并把路径填到里面就行了。
run_test.py

def get_args_parser():
    parser = argparse.ArgumentParser('Set parameters for P2PNet evaluation', add_help=False)
    
    # * Backbone
    parser.add_argument('--backbone', default='vgg16_bn', type=str,
                        help="name of the convolutional backbone to use")

    parser.add_argument('--row', default=2, type=int,
                        help="row number of anchor points")
    parser.add_argument('--line', default=2, type=int,
                        help="line number of anchor points")

    parser.add_argument('--output_dir', default='vis',#添加输出位置定位的路径文件夹
                        help='path where to save')
    parser.add_argument('--weight_path', default='weights/SHTechA.pth',#添加p2pnet的训练好的模型
                        help='path where the trained weights saved')

    parser.add_argument('--gpu_id', default=0, type=int, help='the gpu used for evaluation')

    return parser

def main(args, debug=False):

    os.environ["CUDA_VISIBLE_DEVICES"] = '{}'.format(args.gpu_id)

    print(args)
    device = torch.device('cuda')
    # get the P2PNet
    model = build_model(args)
    # move to GPU
    model.to(device)
    # load trained model
    if args.weight_path is not None:
        checkpoint = torch.load(args.weight_path, map_location='cpu')
        model.load_state_dict(checkpoint['model'])
    # convert to eval mode
    model.eval()
    # create the pre-processing transform
    transform = standard_transforms.Compose([
        standard_transforms.ToTensor(), 
        standard_transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    # set your image path here
    img_path = "./vis/demo1.jpg"#测试数据图片
    # load the images
    img_raw = Image.open(img_path).convert('RGB')
    # round the size
    width, height = img_raw.size
    new_width = width // 128 * 128
    new_height = height // 128 * 128
    img_raw = img_raw.resize((new_width, new_height), Image.ANTIALIAS)
    # pre-proccessing
    img = transform(img_raw)

这里只需要添加自己的输出路径、权重路径、和测试数据的路径就行了。
注意如果出现下面问题

这是版本不统一的问题，只需要将misc.py里的if语句注释掉就行了，如下图

训练：
如果要训练自己的数据集，那需要注意，这个程序进行训练需要的是数据图像和标注位置的txt文件，注意是txt文件，不是公共数据集的mat或csv文件。

import os
import random
import torch
import numpy as np
from torch.utils.data import Dataset
from PIL import Image
import cv2
import glob
import scipy.io as io

class SHHA(Dataset):
    def __init__(self, data_root, transform=None, train=False, patch=False, flip=False):
        self.root_path = data_root
        self.train_lists = "shanghai_tech_part_a_train.list"#训练集路径
        self.eval_list = "shanghai_tech_part_a_test.list"#测试集路径
        # there may exist multiple list files
        self.img_list_file = self.train_lists.split(',')
        if train:
            self.img_list_file = self.train_lists.split(',')
        else:
            self.img_list_file = self.eval_list.split(',')

        self.img_map = {}

这上面的训练集和测试集的文件路径，其实就将训练图片和其标注文件的路径放在一个txt文件中，最后把.txt后缀名改成.list。这个txt文件里的内容是：
图片1路径+空格+ 其对应的是标注文件txt文件的路径
图片2路径+空格+ 其对应的是标注文件txt文件的路径
…
然后改变train.py中的路径

# dataset parameters
    parser.add_argument('--dataset_file', default='SHHA')
    parser.add_argument('--data_root', default='./new_public_density_data',#你放这个程序文件夹的位置
                        help='path where the dataset is')
    
    parser.add_argument('--output_dir', default='./log',#需要自己建立文件夹
                        help='path where to save, empty for no saving')
    parser.add_argument('--checkpoints_dir', default='./ckpt',#需要自己建立文件夹，用来存放训练后的权重文件
                        help='path where to save checkpoints, empty for no saving')
    parser.add_argument('--tensorboard_dir', default='./runs',
                        help='path where to save, empty for no saving')

    parser.add_argument('--seed', default=42, type=int)
    parser.add_argument('--resume', default='', help='resume from checkpoint')
    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                        help='start epoch')
    parser.add_argument('--eval', action='store_true')
    parser.add_argument('--num_workers', default=8, type=int)
    parser.add_argument('--eval_freq', default=5, type=int,
                        help='frequency of evaluation, default setting is evaluating in every 5 epoch')
    parser.add_argument('--gpu_id', default=0, type=int, help='the gpu used for training')

    return parser

def main(args):
    os.environ["CUDA_VISIBLE_DEVICES"] = '{}'.format(args.gpu_id)
    # create the logging file
    run_log_name = os.path.join(args.output_dir, 'run_log.txt')#这个run_log.txt需要自己建立
    with open(run_log_name, "w") as log_file:
        log_file.write('Eval Log %s\n' % time.strftime("%c"))

改变上面的程序文件夹存放的路径，添加训练好的模型存放的文件夹路径，和最后在自己建立的output_dir文件夹中建立run_log.txt文件，就可以进行训练了。
以上就是我们在复现的过程中需要调整的信息了。
如果有哪些小伙伴出现其他问题欢迎在评论区留言。