通常,corwding counter任务中数据集只有图像和对应的标注文件,标注文件中为每个人的坐标(一个人对应一个坐标)。如果需要采用density map进行训练的话,需要生成对应的density map。
本代码用于将mat文件生成对应的density map;最终用.h5文件保存;
import h5py
import scipy.io as io
import PIL.Image as Image
import numpy as np
import os
import glob
from matplotlib import pyplot as plt
from scipy.ndimage.filters import gaussian_filter
from matplotlib import cm as CM
from image import *
# root is the path to ShanghaiTech dataset
root=''
part_B_train = os.path.join(root,'justPoint','images')
# part_B_test = os.path.join(root,'justPoint','images')
path_sets = [part_B_train]
img_paths = []
for path in path_sets:
for img_path in glob.glob(os.path.join(path, '*.jpg')):
img_paths.append(img_path)
for img_path in img_paths:
print(img_path)
mat = io.loadmat(img_path.replace('.jpg','.mat').replace('images','ground_truth').replace('IMG_','GT_IMG_'))
img= plt.imread(img_path)
k = np.zeros((img.shape[0],img.shape[1]))
gt = mat["image_info"][0,0][0,0][0]
for i in range(0,len(gt)):
if int(gt[i][1])<img.shape[0] and int(gt[i][0])<img.shape[1]:
k[int(gt[i][1]),int(gt[i][0])]=1
k = gaussian_filter(k,15)
with h5py.File(img_path.replace('.jpg','.h5').replace('images','ground_truth'), 'w') as hf:
hf['density'] = k