# -*- coding:utf-8 -*-
import os
import glob
from PIL import Image
from scipy import io
import h5py
import numpy as np
from .gaussian_filter import adaptive_gaussian, uniform_gaussian
from tqdm import tqdm
import math
[docs]class PreProcess:
"""generate density map
Args:
root(str) :the root of dataset, only support shtu_a, shtu_b, ucf_qnrf and ucf_cc_50 now.
For ShanghaiTech part A and part B, the root should be upper dir over "part_A_final" or
"part_B_final". And for UCF QNRF and UCF CC 50, the root should be upper dir over "Train" and "Test"
name(str, optional): the name of dataset, must be in ["shtu_a", "shtu_b", "ucf_qnrf", "ucf_cc"]. default: shtu_a
"""
def __init__(self, root, name="shtu_a"):
self.root = root
self.name = name
if self.name not in ["shtu_a", "shtu_b", "ucf_qnrf", "ucf_cc"]:
print("the name should be 'shtu_a', 'shtu_b', 'ucf_qnrf' or 'ucf_cc'")
raise ValueError
self.path_set = {
"shtu_a": [
os.path.join(root, 'part_A_final/train_data', 'images'),
os.path.join(root, 'part_A_final/test_data', 'images')
],
"shtu_b": [
os.path.join(root, 'part_B_final/train_data', 'images'),
os.path.join(root, 'part_B_final/test_data', 'images')
],
"ucf_qnrf": [
os.path.join(root, 'Train'),
os.path.join(root, 'Test')
],
"ucf_cc": [root]
}[name]
[docs] def process(self, mode="uniform", sigma=4, radius=7):
"""generate density map and save.
Args:
mode(str, optional): the way to generate gaussian filter. "uniform": uniform sigma and radius of filter, the
suggestion from C-3-Framework is sigma 15 and radius 7 as the default. "adaptive_kdtree": use adaptive
gaussian kernel with kdtree. "adaptive_voronio": use adaptive gaussian kernel with voronio map. defalut:
uniform.
sigma(int, optional): the sigma of gaussian filter. default: 4.
radius(int, optional): the radius of gaussian area. default: 7.
"""
print("begin to preprocess...")
for path in self.path_set:
img_dir = glob.glob(os.path.join(path, '*.jpg'))
progress_bar = tqdm(total=len(img_dir))
for img_path in img_dir:
img = np.asarray(Image.open(img_path))
if self.name in ["shtu_a", "shtu_b"]:
mat = io.loadmat(img_path.replace('.jpg', '.mat').replace('images/IMG_', 'ground_truth/GT_IMG_'))
gt = mat["image_info"][0, 0][0, 0][0]
save_dir = img_path.replace('.jpg', '.h5').replace('images', 'ground_truth')
else:
mat = io.loadmat(img_path.replace('.jpg', '_ann.mat'))
gt = mat["annPoints"]
save_dir = img_path.replace('.jpg', '.h5')
h, w = img.shape[:2]
density = np.zeros((h, w))
for x, y in gt:
if y < h and x < w:
density[math.floor(y), math.floor(x)] = 1
if mode == "uniform":
density = uniform_gaussian(density, sigma=sigma, radius=radius)
elif mode == "adaptive":
density = adaptive_gaussian(density, mode=mode)
with h5py.File(save_dir, 'r+') as hf:
if 'density' in hf:
hf['density'][...] = density
else:
hf['density'] = density
progress_bar.update(1)
progress_bar.close()