1.BL模型是什么
BL指的是Bayesian Loss,是论文《Bayesian Loss for Crowd Count Estimation with Point Supervision》中提出来的一种损失函数,可用于密度图计数模型。
论文中提到,传统方法需要从点标注生成全图“伪真值标签”,论文中则放弃了这个带来显著噪声和虚假信号的做法,改为从估计的概率密度图上进行期望计算,直接与“真值点标注”进行回归估计,如下图所示,下图是论文作者发在知乎上的博客中的图,对于Bayesian Loss的具体介绍,可见原作者发的文章:https://zhuanlan.zhihu.com/p/127956794,
官方提供的源代码:https://github.com/ZhihengCV/Baysian-Crowd-Counting.
本博客主要介绍数据预处理阶段的一些方法,至于一些具体的神经网络结构等会在后续的博客中介绍。
2 ShanghaiTech数据集介绍
shanghaiTech数据集是一个用于行人计数的数据集,里面包含了大量的人群图片,它由两部分组成:part_A_final文件夹 和 part_B_final文件夹,它们内部都由 train_data 和 test_data 两部分组成,意思为训练集和测试集,对于train_data 和 test_data 其中包括 images 和 ground_truth,images储存着图片文件、ground_truth里储存着对应的标注文件,一张图片对应一个标注文件。
shanghaiTech结构如下图所示:
为了方便后续的处理,本文模仿UCF数据集的结构,将shanghaiTech数据集分为了PA和PB两部分,每部分的train、val、test文件夹存放的是训练集验证集和测试集,训练集取原train_data的所有数据,验证集取原train_data的后10%数据,测试集取原test_data的所有数据,最后本文不再像原shanghaiTech数据集那样把图片和标注文件分开两个images和ground_truth文件夹分开放,而是把图片和标注文件放到了一个文件夹中,具体结构如下图所示:
3 官方数据预处理代码介绍
官方代码提供的preprocess_dataset.py是对UCF数据集进行处理的,它会将UCF数据集分为train、val、test三部分,并且三部分生成的标注文件的结构有所不同,在制作自己的数据集时需要注意,代码已经写的比较清晰了,我再对一些地方进行说明:1.train、val、test三部分的处理有所不同。2.代码会对分辨率过大的图像进行缩放。3.官方将训练集和验证集的图片名放到了相对应的txt文件,在分区时会读取txt文件,preprocess_dataset.py代码如下:
from scipy.io import loadmat
from PIL import Image
import numpy as np
import os
from glob import glob
import cv2
import argparse
def cal_new_size(im_h, im_w, min_size, max_size):
if im_h < im_w:
if im_h < min_size:
ratio = 1.0 * min_size / im_h
im_h = min_size
im_w = round(im_w*ratio)
elif im_h > max_size:
ratio = 1.0 * max_size / im_h
im_h = max_size
im_w = round(im_w*ratio)
else:
ratio = 1.0
else:
if im_w < min_size:
ratio = 1.0 * min_size / im_w
im_w = min_size
im_h = round(im_h*ratio)
elif im_w > max_size:
ratio = 1.0 * max_size / im_w
im_w = max_size
im_h = round(im_h*ratio)
else:
ratio = 1.0
return im_h, im_w, ratio
def find_dis(point):
square = np.sum(point*points, axis=1)
dis = np.sqrt(np.maximum(square[:, None] - 2*np.matmul(point, point.T) + square[None, :], 0.0))
dis = np.mean(np.partition(dis, 3, axis=1)[:, 1:4], axis=1, keepdims=True)
return dis
def generate_data(im_path):
im = Image.open(im_path)
im_w, im_h = im.size
mat_path = im_path.replace('.jpg', '_ann.mat')
points = loadmat(mat_path)['annPoints'].astype(np.float32)
print(points)
idx_mask = (points[:, 0] >= 0) * (points[:, 0] <= im_w) * (points[:, 1] >= 0) * (points[:, 1] <= im_h)
points = points[idx_mask]
im_h, im_w, rr = cal_new_size(im_h, im_w, min_size, max_size)
im = np.array(im)
if rr != 1.0:
im = cv2.resize(np.array(im), (im_w, im_h), cv2.INTER_CUBIC)
points = points * rr
return Image.fromarray(im), points
def parse_args():
parser = argparse.ArgumentParser(description='Test ')
parser.add_argument('--origin-dir', default=r'E:\Project\Net\CSRNet-pytorch-master\CSRNet-pytorch-master\Shanghai\part_A_final',
help='original data directory')
parser.add_argument('--data-dir', default=r'E:\Project\Net\UCF-QNRF_ECCV18 (2)\UCF-QNRF_ECCV18\UCF-Train-Val-Test',
help='processed data directory')
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
save_dir = args.data_dir
min_size = 512
max_size = 2048
for phase in ['Train', 'Test']:
sub_dir = os.path.join(args.origin_dir, phase)
if phase == 'Train':
sub_phase_list = ['train', 'val']
for sub_phase in sub_phase_list:
sub_save_dir = os.path.join(save_dir, sub_phase)
if not os.path.exists(sub_save_dir):
os.makedirs(sub_save_dir)
with open('{}.txt'.format(sub_phase)) as f:
for i in f:
im_path = os.path.join(sub_dir, i.strip())
name = os.path.basename(im_path)
print(name)
im, points = generate_data(im_path)
if sub_phase == 'train':
dis = find_dis(points)
points = np.concatenate((points, dis), axis=1)
im_save_path = os.path.join(sub_save_dir, name)
im.save(im_save_path)
gd_save_path = im_save_path.replace('jpg', 'npy')
np.save(gd_save_path, points)
print("im:"+str(im))
# print("points:" + points)
else:
sub_save_dir = os.path.join(save_dir, 'test')
if not os.path.exists(sub_save_dir):
os.makedirs(sub_save_dir)
im_list = glob(os.path.join(sub_dir, '*jpg'))
for im_path in im_list:
name = os.path.basename(im_path)
print(name)
im, points = generate_data(im_path)
im_save_path = os.path.join(sub_save_dir, name)
im.save(im_save_path)
gd_save_path = im_save_path.replace('jpg', 'npy')
np.save(gd_save_path, points)
4 改写官方预处理代码实现预处理ShanghaiTech数据集
改写后的代码如下,详细的解释都写在了注释中,主要做了以下几点改进:
1.新建了Imgpath函数,实现了传入文件地址,返回地址路径下的.jpg图片文件地址,没有像官方那样把训练集验证集的图片名放到txt里。
2.生成的.npy文件会跟缩放(放大)后图片和mat文件在同一路径下,没有生成新的文件夹来存放数据。
3.参考CSRNet的官方源码(https://github.com/leeyeehoo/CSRNet-pytorch)中的make_dataset.py得到了读取ShanghaiTech中的mat标注文件的代码。
from scipy.io import loadmat
from PIL import Image
import numpy as np
import os
import cv2
import argparse
def cal_new_size(im_h, im_w, min_size, max_size):
if im_h < im_w:
if im_h < min_size:
ratio = 1.0 * min_size / im_h
im_h = min_size
im_w = round(im_w*ratio)
elif im_h > max_size:
ratio = 1.0 * max_size / im_h
im_h = max_size
im_w = round(im_w*ratio)
else:
ratio = 1.0
else:
if im_w < min_size:
ratio = 1.0 * min_size / im_w
im_w = min_size
im_h = round(im_h*ratio)
elif im_w > max_size:
ratio = 1.0 * max_size / im_w
im_w = max_size
im_h = round(im_h*ratio)
else:
ratio = 1.0
return im_h, im_w, ratio
def find_dis(point):
square = np.sum(point*points, axis=1)
dis = np.sqrt(np.maximum(square[:, None] - 2*np.matmul(point, point.T) + square[None, :], 0.0))
dis = np.mean(np.partition(dis, 3, axis=1)[:, 1:4], axis=1, keepdims=True)
return dis
def generate_data(im_path):
im = Image.open(im_path)
im_w, im_h = im.size
# 获取标注点坐标
mat_path = im_path.replace('.jpg','.mat').replace('images','ground_truth').replace('IMG_','GT_IMG_')
points = loadmat(mat_path)['image_info'][0,0][0,0][0]
# print(points)
idx_mask = (points[:, 0] >= 0) * (points[:, 0] <= im_w) * (points[:, 1] >= 0) * (points[:, 1] <= im_h)
points = points[idx_mask]
# 图像缩放(放大),返回缩放后的图片宽高和缩放比例
im_h, im_w, rr = cal_new_size(im_h, im_w, min_size, max_size)
im = np.array(im)
# 根据缩放比例处理点坐标信息
if rr != 1.0:
im = cv2.resize(np.array(im), (im_w, im_h), cv2.INTER_CUBIC)
points = points * rr
return Image.fromarray(im), points
def Imgpath(path):
# 获取图片文件地址
files = os.listdir(path)
img_paths = []
for filename in files:
portion = os.path.splitext(filename)
if portion[1] == '.jpg':
img_paths.append(os.path.join(sub_save_dir, filename))
return img_paths
def parse_args():
parser = argparse.ArgumentParser(description='Test ')
parser.add_argument('--origin-dir', default=r'.\Shanghai\part_B_final',
help='original data directory')
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
min_size = 512
max_size = 2048
for phase in ['Train', 'Test']:
# 处理训练集和验证集
if phase == 'Train':
sub_phase_list = ['train', 'val']
for sub_phase in sub_phase_list:
sub_save_dir = os.path.join(args.origin_dir, sub_phase)
# 获取图片文件地址
img_paths = Imgpath(sub_save_dir)
# 遍历图片文件地址
for im_path in img_paths:
# 获取图片名
name = os.path.basename(im_path)
print(sub_phase + ": " + name)
# 生成图片文件im和点坐标points
im, points = generate_data(im_path)
# 如果是训练集则做额外的处理
if sub_phase == 'train':
dis = find_dis(points)
points = np.concatenate((points, dis), axis=1)
# 保存缩放后的图片和点坐标
im_save_path = os.path.join(sub_save_dir, name)
im.save(im_save_path)
gd_save_path = im_save_path.replace('jpg', 'npy')
np.save(gd_save_path, points)
# 处理测试集
else:
sub_save_dir = os.path.join(args.origin_dir, 'test')
img_paths = Imgpath(sub_save_dir)
for im_path in img_paths:
name = os.path.basename(im_path)
print("test: " + name)
# print(name)
im, points = generate_data(im_path)
im_save_path = os.path.join(sub_save_dir, name)
im.save(im_save_path)
gd_save_path = im_save_path.replace('jpg', 'npy')
np.save(gd_save_path, points)
处理后文件结构如下所示:
其中一个训练集内部的结构如下(其它的训练集、测试集和验证集也类似,mat、npy和jpg文件在同一路径下),所示:
5 小结
本文通过改写BL模型官方源代码中的preprocess_dataset.py文件,实现了将ShanghaiTech数据集处理为BL模型所需的数据格式,其实并不难,本质上就是提取出关键点坐标进行处理,所以如果想要使用自己的关键点标注的数据集来跑BL模型,仅需将我的代码中提取关键点坐标的代码进行改写即可。