关于shanghaiTech的介绍,可以看我之前的写的博客(https://blog.csdn.net/johnhamtom/article/details/104481186?spm=1001.2014.3001.5501)
这里采用的文件结构如下图所示,其中images存放图片,ground_truth存放对应的标注文件:
参考CSRNet的官方源码(https://github.com/leeyeehoo/CSRNet-pytorch)中的make_dataset.py可以得知,读取mat标注文件并获取其点坐标的代码如下:
mat = io.loadmat(img_path.replace('.jpg','.mat').replace('images','ground_truth').replace('IMG_','GT_IMG_'))
gt = mat["image_info"][0,0][0,0][0] #gt为用列表存储的点坐标对
知道如何获取mat文件的点坐标后,直接批量读取mat文件,并在读取的时候记录关键点数量和图片名称即可,一些说明写在了注释中,代码如下:
import scipy.io as io
import glob
from image import *
import pandas as pd
# 设置根目录地址
root = r'.\Shanghai'
# 根据根目录地址合成图片地址
part_A_train = os.path.join(root, 'part_A_final', 'train_data', 'train_data', 'images')
part_A_val = os.path.join(root, 'part_A_final', 'train_data', 'val_data', 'images')
part_A_test = os.path.join(root, 'part_A_final', 'test_data', 'images')
part_B_train = os.path.join(root, 'part_B_final', 'train_data', 'train_data', 'images')
part_B_val = os.path.join(root, 'part_B_final', 'train_data', 'val_data', 'images')
part_B_test = os.path.join(root, 'part_B_final', 'test_data', 'images')
path_sets = [part_A_train, part_A_val, part_A_test, part_B_train, part_B_val, part_B_test]
save_path = ['part_A_train', 'part_A_val', 'part_A_test', 'part_B_train', 'part_B_val', 'part_B_test']
index_path_sets = 0
for name in save_path:
# 读取图片路径
img_paths = []
for img_path in glob.glob(os.path.join(path_sets[index_path_sets], '*.jpg')):
img_paths.append(img_path)
index_path_sets = index_path_sets + 1
df = pd.DataFrame(columns=['ImgName', 'Num']) # 将标注框数量储存在CSV中
#文件名
filename = name +"_num"
i=0
with open(r".\Shanghai\annotations/" + filename + ".txt" , "w") as f:
for img_path in img_paths:
mat = io.loadmat(img_path.replace('.jpg','.mat').replace('images','ground_truth').replace('IMG_','GT_IMG_'))
gt = mat["image_info"][0, 0][0, 0][0]
# len(gt) 即为标注点数量
# os.path.basename(img_path) 提取出图片文件名
new = [os.path.basename(img_path), len(gt)]
df.loc[i] = new
print(new)
i=i+1
f.write(os.path.basename(img_path) + " " + str(len(gt)))
f.write("\n")
df.to_csv(r".\Shanghai\annotations/" + filename + ".csv", index=False) # CSV文件路径
结果如下所示: