def get_blob(self, filename):
image_name, _ = os.path.splitext(filename)
blob = dict()
blob['image_name'] = image_name
if self.image_channel == 1:
image = cv2.imread(join(self.image_path, filename), 0)
elif self.image_channel == 3:
image = cv2.imread(join(self.image_path, filename), 1)
else:
raise Exception('invalid number of image channels')
density_map = pandas.read_csv(join(self.density_map_path,
image_name + '.csv'),
sep=',',
header=None).values
if image.shape[0] != density_map.shape[0] or image.shape[
1] != density_map.shape[1]:
raise Exception('density map size mismatch.')
density_map = self.downsample(density_map, DOWNSAMPLE)
if self.roi_path is not None:
if DOWNSAMPLE == 1:
roi = self.load_roi(
join(self.roi_path, image_name + '_roi.mat'))
elif DOWNSAMPLE == 4:
roi = self.load_roi(
join(self.roi_path, image_name + '_roi_fourth_size.mat'))
elif DOWNSAMPLE == 8:
roi = self.load_roi(
join(self.roi_path, image_name + '_roi_eighth_size.mat'))
else:
raise Exception('no suitable RoI file available')
else:
roi = None
image = self.reshape_data(image)
density_map = self.reshape_data(density_map)
if roi is not None:
roi = self.reshape_data(roi)
if roi.shape[2] != density_map.shape[2] or roi.shape[
3] != density_map.shape[3]:
raise Exception('RoI size mismatch')
blob['image'] = ndarray_to_tensor(image, is_cuda=False)
blob['density'] = ndarray_to_tensor(density_map, is_cuda=False)
if roi is not None:
blob['roi'] = ndarray_to_tensor(roi, is_cuda=False)
else:
blob['roi'] = None
if self.is_label:
blob['label'] = ndarray_to_tensor(self.get_label(blob['density']),
is_cuda=False)
else:
blob['label'] = None
return blob
def read_blob(self, filename):
image_name, _ = os.path.splitext(filename)
blob = dict()
blob['image_name'] = image_name
# read image
image = cv2.imread(join(self.image_path, filename), 1)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
density_map = pandas.read_csv(join(self.density_map_path, image_name + '.csv'), sep=',', header=None).values
if image.shape[0] != density_map.shape[0] or image.shape[1] != density_map.shape[1]:
raise Exception('density map size mismatch.')
density_map = self.downsample(density_map, DOWNSAMPLE)
if self.roi_path is not None:
# if DOWNSAMPLE == 1:
# roi = self.load_roi(join(self.roi_path, image_name + '_roi.mat'))
# elif DOWNSAMPLE == 4:
# roi = self.load_roi(join(self.roi_path, image_name + '_roi_fourth_size.mat'))
# elif DOWNSAMPLE == 8:
# roi = self.load_roi(join(self.roi_path, image_name + '_roi_eighth_size.mat'))
# else:
# raise Exception('no suitable RoI file available')
roi = self.load_roi(join(self.roi_path, image_name + '_roi.mat'))
else:
roi = None
# image = self.reshape_data(image)
image = self.image2tensor(image)
density_map = self.reshape_data(density_map)
if roi is not None:
roi = self.reshape_data(roi)
if roi.shape[1] != image.shape[1] or roi.shape[2] != image.shape[2]:
raise Exception('RoI size mismatch')
else:
roi = np.ones((1, image.shape[1], image.shape[2]))
if isinstance(image, torch.Tensor):
blob['image'] = image
else:
blob['image'] = ndarray_to_tensor(image, is_cuda=False)
blob['density'] = ndarray_to_tensor(density_map, is_cuda=False)
blob['roi'] = ndarray_to_tensor(roi, is_cuda=False)
if self.is_label:
blob['label'] = ndarray_to_tensor(self.get_label(blob['density']), is_cuda=False)
if self.is_mask:
blob['mask'] = ndarray_to_tensor(self.get_mask(blob['density'], blob['roi']), is_cuda=False)
return blob
def calculate_error(ground_truth, estimate, L=0):
# grid average mean absolute error
# ground_truth Tensor: shape=(1, 1, h, w)
# estimate Tensor: same shape of ground_truth
ground_truth = ndarray_to_tensor(ground_truth, is_cuda=True)
estimate = ndarray_to_tensor(estimate, is_cuda=True)
height = ground_truth.shape[2]
width = ground_truth.shape[3]
times = math.sqrt(math.pow(4, L))
padding_height = int(math.ceil(height / times) * times - height)
padding_width = int(math.ceil(width / times) * times - width)
if padding_height != 0 or padding_width != 0:
m = nn.ZeroPad2d((0, padding_width, 0, padding_height))
ground_truth = m(ground_truth)
estimate = m(estimate)
height = ground_truth.shape[2]
width = ground_truth.shape[3]
m = nn.AdaptiveAvgPool2d(int(times))
ground_truth = m(ground_truth) * (height / times) * (width / times)
estimate = m(estimate) * (height / times) * (width / times)
game = torch.sum(torch.abs(ground_truth - estimate))
return game.item()
def predict(self, frame):
height = frame.shape[0]
width = frame.shape[1]
process_height = 400
process_width = int(width / height * process_height)
frame = cv2.resize(frame, (process_width, process_height))
# get original size
height = frame.shape[0]
width = frame.shape[1]
reshaped_frame = np.moveaxis(frame, 2, 0).astype(
np.float32) # reshape (h, w, 3) to (3, h, w)
reshaped_frame = reshaped_frame.reshape((1, 3, height, width))
image_data = ndarray_to_tensor(reshaped_frame, is_cuda=True)
estimate_map, _ = self.net(image_data)
estimate_map = estimate_map.data.cpu().numpy()
estimate_count = np.sum(estimate_map)
max_value = np.max(estimate_map)
if max_value > 0:
estimate_prior_normalized = estimate_map[0][0] / np.max(
estimate_map)
else:
estimate_prior_normalized = estimate_map[0][0]
estimate_prior_normalized_bgr = gray_to_bgr(estimate_prior_normalized)
image_estimate_map = cv2.addWeighted(
frame, self.alpha,
cv2.resize(estimate_prior_normalized_bgr, (width, height)),
1 - self.alpha, 0)
estimate_count_text = '%.f' % estimate_count
t_size = cv2.getTextSize(estimate_count_text, cv2.FONT_HERSHEY_PLAIN,
1, 1)[0]
cv2.putText(image_estimate_map, estimate_count_text,
(0, t_size[1] + 4), cv2.FONT_HERSHEY_PLAIN, 1, [0, 0, 255],
1)
return image_estimate_map, estimate_count
def get_mask(self, ground_truth_map):
# ground_truth numpy.ndarray shape=(1, 1, h, w)
# return mask numpy.ndarray shape=(1, 3, h, w)
foreground_mask = np.zeros_like(ground_truth_map)
foreground_mask[ground_truth_map > 0] = 1.0
ground_truth = ndarray_to_tensor(ground_truth_map)
pooled_ground_truth = self.average_pool(ground_truth)
pooled_ground_truth = pooled_ground_truth.data.cpu().numpy()
# start_density = 0
# end_density = self.low_density_flag
# zeros = np.zeros_like(pooled_ground_truth)
# zeros[np.logical_and(pooled_ground_truth > start_density, pooled_ground_truth <= end_density)] = 1
# mask = zeros * foreground_mask
#
# start_density = self.low_density_flag
# end_density = self.high_density_flag
# zeros = np.zeros_like(pooled_ground_truth)
# zeros[np.logical_and(pooled_ground_truth > start_density, pooled_ground_truth <= end_density)] = 1
# mask = np.concatenate((mask, zeros * foreground_mask), axis=1)
#
# start_density = self.high_density_flag
# end_density = sys.maxsize
# zeros = np.zeros_like(pooled_ground_truth)
# zeros[np.logical_and(pooled_ground_truth > start_density, pooled_ground_truth <= end_density)] = 1
# mask = np.concatenate((mask, zeros * foreground_mask), axis=1)
zeros = np.zeros_like(pooled_ground_truth)
zeros[pooled_ground_truth < self.density_flag] = 1
mask = zeros * foreground_mask
zeros = np.zeros_like(pooled_ground_truth)
zeros[pooled_ground_truth >= self.density_flag] = 1
mask = np.concatenate((mask, zeros * foreground_mask), axis=1)
return mask
def __init__(self, data, size=9):
if not isinstance(size, int):
raise Exception('size should be an integer')
if size % 2 == 0:
raise Exception('size should be an odd integer')
self.average_pool = nn.AvgPool2d(size,
stride=1,
padding=int((size - 1) / 2),
count_include_pad=False)
index = 0
number_of_samples = data.get_number_of_samples()
density_list = list()
print('Building mask label.')
for blob in data:
index += 1
ground_truth_data = blob['density_map']
if np.sum(ground_truth_data) == 0:
continue
ground_truth = ndarray_to_tensor(ground_truth_data)
pooled_ground_truth = self.average_pool(ground_truth)
pooled_ground_truth = pooled_ground_truth.data.cpu().numpy()
density_list.extend(pooled_ground_truth[pooled_ground_truth > 0])
if index % 100 == 0:
print('Built %6d / %d mask labels.' %
(index, number_of_samples))
self.density_flag = np.percentile(density_list, 75)
print('Finished building mask.')
TIMES = 16
process_height = int(process_height / TIMES) * TIMES
process_width = int(process_width / TIMES) * TIMES
frame = cv2.resize(frame, (process_width, process_height))
# get original size
height = frame.shape[0]
width = frame.shape[1]
reshaped_frame = np.moveaxis(frame, 2, 0).astype(
np.float32) # reshape (h, w, 3) to (3, h, w)
reshaped_frame = reshaped_frame.reshape((1, 3, height, width))
image_data = ndarray_to_tensor(reshaped_frame, is_cuda=True)
estimate_map, _ = net(image_data)
estimate_map = estimate_map.data.cpu().numpy()
estimate_count = np.sum(estimate_map)
max_value = np.max(estimate_map)
if max_value > 0:
estimate_prior_normalized = estimate_map[0][0] / np.max(
estimate_map)
else:
estimate_prior_normalized = estimate_map[0][0]
estimate_prior_normalized_bgr = gray_to_bgr(estimate_prior_normalized)
image_estimate_map = cv2.addWeighted(
def multithread_dataloader(data_config):
# data_config: dict, a dictionay contains several datasets info,
# key is dataset name,
# value is a dict which contains is_preload and is_label and is_mask
data_path = DataPath()
data_dict = dict()
for name in data_config:
this_dataset_flag = data_config[name]
if 'label' in this_dataset_flag:
is_label = this_dataset_flag['label']
else:
is_label = False
if 'mask' in this_dataset_flag:
is_mask = this_dataset_flag['mask']
else:
is_mask = False
if 'shuffle' in this_dataset_flag:
is_shuffle = this_dataset_flag['shuffle']
else:
is_shuffle = False
if 'seed' in this_dataset_flag:
random_seed = this_dataset_flag['seed']
else:
random_seed = None
if 'batch_size' in this_dataset_flag:
batch_size = this_dataset_flag['batch_size']
else:
batch_size = 1
if random_seed is not None:
def worker_init_fn(x):
seed = random_seed + x
np.random.seed(seed)
random.seed(seed)
torch.manual_seed(seed)
return
else:
worker_init_fn = None
path = data_path.get_path(name)
this_data = Data(name,
path['image'],
path['gt'],
roi_path=path['roi'],
is_label=is_label,
is_mask=is_mask)
this_dataloader = DataLoader(this_data,
batch_size=batch_size,
shuffle=is_shuffle,
num_workers=8,
drop_last=False,
worker_init_fn=worker_init_fn)
if is_label:
label_histogram = np.zeros(NUMBER_OF_LABELS)
index = 0
for blob in this_dataloader:
label_histogram[torch.argmax(blob['label'])] += 1
index += 1
if index % 100 == 0:
print('Built %6d of %d labels.' %
(index, this_data.get_number_of_samples()))
print('Completed building %d labels. Label histogram is %s' %
(index, ' '.join([str(i) for i in label_histogram])))
label_weights = 1 - label_histogram / sum(label_histogram)
label_weights = label_weights / sum(label_weights)
else:
label_weights = None
this_dataset_dict = dict()
this_dataset_dict['data'] = this_dataloader
if is_label:
this_dataset_dict['label_weights'] = ndarray_to_tensor(
label_weights, is_cuda=False)
data_dict[name] = this_dataset_dict
return data_dict
def evaluate_model(model_path, data):
net = CrowdCount()
network.load_net(model_path, net)
net.cuda()
net.eval()
build_ssim = SSIM(window_size=11)
game = GridAverageMeanAbsoluteError()
mae = 0.0
mse = 0.0
psnr = 0.0
ssim = 0.0
game_0 = 0.0
game_1 = 0.0
game_2 = 0.0
game_3 = 0.0
index = 0
for blob in data:
image_data = blob['image']
ground_truth_data = blob['density']
roi = blob['roi']
# filename = blob['filename']
if image_data.shape[0] != 1:
raise Exception('invalid image batch size (%d) for evaluation' %
image_data.shape[0])
estimate_map, _ = net(image_data, roi=roi)
ground_truth_data = ground_truth_data.data.cpu().numpy()
density_map = estimate_map.data.cpu().numpy()
ground_truth_count = np.sum(ground_truth_data)
estimate_count = np.sum(density_map)
mae += abs(ground_truth_count - estimate_count)
mse += (ground_truth_count - estimate_count)**2
psnr += build_psnr(ground_truth_data, density_map)
ssim += build_ssim(ndarray_to_tensor(ground_truth_data),
ndarray_to_tensor(density_map)).item()
game_0 += game.calculate_error(ground_truth_data, density_map, 0)
game_1 += game.calculate_error(ground_truth_data, density_map, 1)
game_2 += game.calculate_error(ground_truth_data, density_map, 2)
game_3 += game.calculate_error(ground_truth_data, density_map, 3)
index += 1
result_dict = dict()
result_dict['name'] = os.path.basename(model_path)
result_dict['number'] = int(index)
result_dict['mae'] = float(mae / index)
result_dict['mse'] = float(np.sqrt(mse / index))
result_dict['psnr'] = float(psnr / index)
result_dict['ssim'] = float(ssim / index)
result_dict['game_0'] = float(game_0 / index)
result_dict['game_1'] = float(game_1 / index)
result_dict['game_2'] = float(game_2 / index)
result_dict['game_3'] = float(game_3 / index)
return result_dict
def forward(self, im_data, roi=None):
with torch.no_grad():
x_prior = self.prior(im_data)
flag = torch.argmax(x_prior, dim=1, keepdim=True)
background_mask = (flag == 0).to(torch.float32)
foreground_mask = 1 - background_mask
resized_foreground_mask = functional.interpolate(1 -
background_mask,
scale_factor=8.0,
mode='nearest')
# mask = None
# for i in range(1, x_prior.shape[1]):
# if mask is None:
# mask = (flag == i).to(torch.float32)
# else:
# mask = torch.cat((mask, (flag == i).to(torch.float32)), dim=1)
low_density_mask = (flag == 1).to(torch.float32)
high_density_mask = (flag == 2).to(torch.float32)
low_density_mask_np = low_density_mask.data.cpu().numpy()
high_density_mask_np = high_density_mask.data.cpu().numpy()
dilate_kernel = np.ones((2, 2))
dilated_low_density_mask_list = list()
dilated_high_density_mask_list = list()
for batch in range(low_density_mask_np.shape[0]):
this_mask = cv2.dilate(low_density_mask_np[batch][0],
dilate_kernel,
iterations=1)
dilated_low_density_mask_list.append(
ndarray_to_tensor(
this_mask.reshape((1, 1, this_mask.shape[0],
this_mask.shape[1]))).cuda())
this_mask = cv2.dilate(high_density_mask_np[batch][0],
dilate_kernel,
iterations=1)
dilated_high_density_mask_list.append(
ndarray_to_tensor(
this_mask.reshape((1, 1, this_mask.shape[0],
this_mask.shape[1]))).cuda())
dilated_low_density_mask = torch.cat(dilated_low_density_mask_list,
dim=0)
dilated_high_density_mask = torch.cat(
dilated_high_density_mask_list, dim=0)
overlap_mask = (
low_density_mask * dilated_high_density_mask +
high_density_mask * dilated_low_density_mask) * foreground_mask
new_low_density_mask = low_density_mask * (
1 - overlap_mask) * foreground_mask
new_high_density_mask = high_density_mask * (
1 - overlap_mask) * foreground_mask
x1 = self.vgg16(im_data * resized_foreground_mask)
maps = self.map(x1)
scales = self.scale(x1) + 1
low_density_map, high_density_map = torch.chunk(maps, 2, dim=1)
low_density_scale, high_density_scale = torch.chunk(scales, 2, dim=1)
scaled_low_density_map = low_density_map * low_density_scale * new_low_density_mask
scaled_high_density_map = high_density_map * high_density_scale * new_high_density_mask
scaled_overlap_map = (
low_density_map * low_density_scale * overlap_mask +
high_density_map * high_density_scale * overlap_mask) / 2
scaled_maps = torch.cat((scaled_low_density_map,
scaled_high_density_map, scaled_overlap_map),
dim=1)
density_map = torch.sum(scaled_maps, 1, keepdim=True)
if roi is not None:
density_map = density_map * roi
visual_dict = dict()
# visual_dict['score'] = x_score_maps
# visual_dict['class'] = x_class_maps
visual_dict['density'] = density_map
visual_dict['raw_maps'] = maps
visual_dict['scaled_maps'] = scaled_maps
visual_dict['masks'] = torch.cat(
(background_mask, new_low_density_mask, new_high_density_mask,
overlap_mask),
dim=1)
return density_map, foreground_mask, visual_dict