def test(train_test_unit, out_dir_root):
output_dir = osp.join(out_dir_root, train_test_unit.metadata['name'])
mkdir_if_missing(output_dir)
sys.stdout = Logger(osp.join(output_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
val_path = train_test_unit.test_dir_img
val_gt_path = train_test_unit.test_dir_den
if not args.resume:
pretrained_model = osp.join(output_dir, 'best_model.h5')
else:
if args.resume[-3:] == '.h5':
pretrained_model = args.resume
else:
resume_dir = osp.join(args.resume,
train_test_unit.metadata['name'])
pretrained_model = osp.join(resume_dir, 'best_model.h5')
print("Using {} for testing.".format(pretrained_model))
data_loader = ImageDataLoader(val_path,
val_gt_path,
shuffle=False,
batch_size=1)
mae, mse = evaluate_model(pretrained_model,
data_loader,
save_test_results=args.save_plots,
plot_save_dir=osp.join(output_dir,
'plot-results-test/'))
print("MAE: {0:.4f}, MSE: {1:.4f}".format(mae, mse))
def evaluate_model(trained_model,
data_loader,
save_test_results=False,
plot_save_dir="/tmp/"):
net = CrowdCounter()
network.load_net(trained_model, net)
net.cuda()
net.eval()
mae = 0.0
mse = 0.0
for blob in data_loader:
im_data = blob['data']
gt_data = blob['gt_density']
idx_data = blob['idx']
im_data_norm = im_data / 255.0
density_map = net(im_data_norm, gt_data)
density_map = density_map.data.cpu().numpy()
gt_count = np.sum(gt_data)
et_count = np.sum(density_map)
mae += abs(gt_count - et_count)
mse += ((gt_count - et_count) * (gt_count - et_count))
if save_test_results:
print("Plotting results")
mkdir_if_missing(plot_save_dir)
save_results(im_data, gt_data, density_map, idx_data,
plot_save_dir)
mae = mae / data_loader.get_num_samples()
mse = np.sqrt(mse / data_loader.get_num_samples())
return mae, mse
def _create_original_density_maps(self, force_create_den_maps):
all_density_dirs_exist = osp.exists(self.ori_dir_partA_train_den)
all_density_dirs_exist = all_density_dirs_exist and osp.exists(
self.ori_dir_partA_test_den)
all_density_dirs_exist = all_density_dirs_exist and osp.exists(
self.ori_dir_partB_train_den)
all_density_dirs_exist = all_density_dirs_exist and osp.exists(
self.ori_dir_partB_test_den)
if not all_density_dirs_exist:
mkdir_if_missing(self.ori_dir_partA_train_den)
mkdir_if_missing(self.ori_dir_partA_test_den)
mkdir_if_missing(self.ori_dir_partB_train_den)
mkdir_if_missing(self.ori_dir_partB_test_den)
elif not force_create_den_maps:
return
create_density_map(self.ori_dir_partA_train_img,
self.ori_dir_partA_train_lab,
self.ori_dir_partA_train_den,
self.ori_dir_partA_train_fac,
mode=self.metadata['gt_mode'])
create_density_map(self.ori_dir_partA_test_img,
self.ori_dir_partA_test_lab,
self.ori_dir_partA_test_den,
self.ori_dir_partA_test_fac,
mode=self.metadata['gt_mode'])
create_density_map(self.ori_dir_partB_train_img,
self.ori_dir_partB_train_lab,
self.ori_dir_partB_train_den,
self.ori_dir_partB_train_fac,
mode=self.metadata['gt_mode'])
create_density_map(self.ori_dir_partB_test_img,
self.ori_dir_partB_test_lab,
self.ori_dir_partB_test_den,
self.ori_dir_partB_test_fac,
mode=self.metadata['gt_mode'])
def save_results(img, gt_density_map, et_density_map, idx, output_dir):
idx = idx[0]
img = img[0, 0]
gt_density_map = np.array(gt_density_map[0, 0])
et_density_map = et_density_map[0, 0]
gt_count = np.sum(gt_density_map)
et_count = np.sum(et_density_map)
maxi = gt_density_map.max()
raw_dir = os.path.join(output_dir, 'raw_plots')
mkdir_if_missing(raw_dir)
if maxi != 0:
gt_density_map = gt_density_map * (255. / maxi)
et_density_map = et_density_map * (255. / maxi)
if gt_density_map.shape[1] != img.shape[1]:
gt_density_map = cv2.resize(gt_density_map,
(img.shape[1], img.shape[0]))
et_density_map = cv2.resize(et_density_map,
(img.shape[1], img.shape[0]))
img_file_name = os.path.join(output_dir, str(idx) + ".jpg")
raw_img_file_name = os.path.join(raw_dir, str(idx) + ".jpg")
cv2.imwrite(raw_img_file_name.replace('.jpg', '_gt.jpg'),
gt_density_map / np.max(gt_density_map) * 255)
cv2.imwrite(raw_img_file_name.replace('.jpg', '_et.jpg'),
et_density_map / np.max(et_density_map) * 255)
fig = plt.figure(figsize=(30, 20))
a = fig.add_subplot(1, 3, 1)
plt.imshow(img, cmap='gray')
a.set_title('input')
plt.axis('off')
a = fig.add_subplot(1, 3, 2)
plt.imshow(gt_density_map)
a.set_title('ground thruth {:.2f}'.format(gt_count))
plt.axis('off')
a = fig.add_subplot(1, 3, 3)
plt.imshow(et_density_map)
a.set_title('estimated {:.2f}'.format(et_count))
plt.axis('off')
#fig.savefig(img_file_name, bbox_inches='tight')
fig.clf()
plt.close()
del a
gc.collect()
def evaluate_model(trained_model,
data_loader,
epoch=0,
save_test_results=False,
plot_save_dir="/tmp/",
den_factor=1e3):
net = CrowdCounter()
network.load_net(trained_model, net)
net.cuda()
net.eval()
mae = 0.0
mse = 0.0
for blob in data_loader:
im_data = blob['data']
gt_data = blob['gt_density']
idx_data = blob['idx']
new_shape = blob['new_shape']
orig_shape = blob['orig_shape']
im_data_norm = im_data / 127.5 - 1. #normalize between -1 and 1
gt_data = gt_data * den_factor
density_map = net(im_data_norm, epoch=epoch)
density_map = density_map.data.cpu().numpy()
density_map /= den_factor
gt_data /= den_factor
im_data, gt_data = data_loader.recontruct_test(im_data, gt_data,
orig_shape, new_shape)
_, density_map = data_loader.recontruct_test(im_data_norm, density_map,
orig_shape, new_shape)
gt_count = np.sum(gt_data)
et_count = np.sum(density_map)
print("image {} gt {:.3f} es {:.3f}".format(idx_data[0], gt_count,
et_count))
mae += abs(gt_count - et_count)
mse += ((gt_count - et_count) * (gt_count - et_count))
if save_test_results:
print("Plotting results")
mkdir_if_missing(plot_save_dir)
utils.save_results(im_data, gt_data, density_map, idx_data,
plot_save_dir)
mae = mae / data_loader.get_num_samples()
mse = np.sqrt(mse / data_loader.get_num_samples())
return mae, mse
def run_face_detector(img_dir_path, den_dir_path, est_dir_path, lab_dir_path):
mkdir_if_missing(est_dir_path)
file_names = os.listdir(img_dir_path)
file_names.sort()
mae, mse = 0, 0
for file_name in file_names:
file_extention = file_name.split('.')[-1]
file_id = file_name[:len(file_name) - len(file_extention)]
file_path = os.path.join(img_dir_path, file_name)
density_map_path = os.path.join(den_dir_path, file_id + 'npy')
label_path = os.path.join(lab_dir_path, file_id + 'json')
img = cv2.imread(file_path)
boxes = detector.detect(img)
save_path = os.path.join(est_dir_path, file_id + 'npy')
np.save(save_path, boxes)
gt = np.load(density_map_path)
gt_count = np.sum(gt)
et_count = boxes.shape[0]
mae += abs(gt_count-et_count)
mse += ((gt_count-et_count)*(gt_count-et_count))
if args.save_plots:
for r in boxes:
cv2.rectangle(img, (r[0],r[1]), (r[2],r[3]), (255,255,0), 1)
points = [[p['y'], p['x']] for p in json.load(open(label_path))]
_, bb_sizes = interpolate_scale(img.shape, points, boxes)
for p, r in zip(points, bb_sizes):
cv2.rectangle(img, (r[0],r[1]), (r[2],r[3]), (255,0, 255), 1)
cv2.circle(img, (int(p[1]), int(p[0])), 1, (255,0, 255), 2)
#img = cv2.resize(img, (0,0), fx=0.5, fy=0.5)
print('Image {}: gt count {:.1f}, est count {}'.format(file_id, gt_count, et_count))
if not os.path.exists(est_dir_path):
os.mkdir(est_dir_path)
save_path = os.path.join(est_dir_path, file_id + 'jpg')
print (save_path)
cv2.imwrite(save_path, img)
mae = mae/len(file_names)
mse = np.sqrt(mse/len(file_names))
print("final mse: {:.3f}, final mae: {:.3f}".format(mse, mae))
def cam(net, blob, density_map, save_dir, method):
mkdir_if_missing(save_dir)
net = net.net
net.eval()
net.zero_grad()
net.cam = True
torch.sum(density_map).backward(retain_graph=True)
gradients = net.get_gradients()
im_data_norm = torch.Tensor(blob['data'] / 127.5 - 1.).to('cuda')
activations = net.get_activations(im_data_norm)
maps = []
for column in range(len(gradients)):
act = activations[column].detach().cpu()
if method == 'grad-cam':
weight = torch.mean(gradients[column], dim=[0, 2, 3]).cpu()
#elif method == 'eigen-cam':
# weight = get_2d_projection(act.data.numpy())
elif method == 'grad-cam++':
weight = weight_grad_cam_pp(act, gradients[column])
elif method == 'layer-cam':
weight = get_layer_cam(act, gradients[column])
if method != 'layer-cam':
for i in range(act.shape[1]):
act[:, i, :, :] *= weight[i]
else:
act = weight
heatmap = torch.mean(act, dim=1).squeeze()
#heatmap = torch.abs(heatmap)
heatmap = np.maximum(heatmap, 0)
heatmap = heatmap.detach().numpy()
heatmap = (heatmap - np.min(heatmap)) / (np.max(heatmap) -
np.min(heatmap))
mean, std = np.mean(heatmap), np.std(heatmap)
#heatmap[heatmap<mean+0.1*std] = 0
maps.append(heatmap)
#max_val = max(list(map(lambda x: np.max(x), maps)))
#min_val = min(list(map(lambda x: np.min(x), maps)))
for column in range(len(gradients)):
heatmap = maps[column]
plot_cam(blob, heatmap, save_dir, column)
def _create_original_density_maps(self, force_create_den_maps):
mkdir_if_missing(self.ori_train_den)
mkdir_if_missing(self.ori_val_den)
mkdir_if_missing(self.ori_test_den)
if len(os.listdir(self.ori_train_den)) != len(
os.listdir(self.ori_train_gt)):
create_density_map(self.ori_train_img,
self.ori_train_lab,
self.ori_train_den,
self.ori_train_fac,
mode=self.metadata["gt_mode"])
if len(os.listdir(self.ori_val_den)) != len(os.listdir(
self.ori_val_gt)):
create_density_map(self.ori_val_img,
self.ori_val_lab,
self.ori_val_den,
self.ori_val_fac,
mode=self.metadata["gt_mode"])
if len(os.listdir(self.ori_test_den)) != len(
os.listdir(self.ori_test_gt)):
create_density_map(self.ori_test_img,
self.ori_test_lab,
self.ori_test_den,
self.ori_test_fac,
mode=self.metadata["gt_mode"])
def _create_labels(self):
mkdir_if_missing(self.ori_dir_partA_train_lab)
mkdir_if_missing(self.ori_dir_partA_test_lab)
mkdir_if_missing(self.ori_dir_partB_train_lab)
mkdir_if_missing(self.ori_dir_partB_test_lab)
# check if number os files is equal
if len(os.listdir(self.ori_dir_partA_train_mat)) != len(os.listdir(self.ori_dir_partA_train_lab)):
self._convert_mat_to_json(self.ori_dir_partA_train_mat, self.ori_dir_partA_train_lab)
if len(os.listdir(self.ori_dir_partA_test_mat)) != len(os.listdir(self.ori_dir_partA_test_lab)):
self._convert_mat_to_json(self.ori_dir_partA_test_mat, self.ori_dir_partA_test_lab)
if len(os.listdir(self.ori_dir_partB_train_mat)) != len(os.listdir(self.ori_dir_partB_train_lab)):
self._convert_mat_to_json(self.ori_dir_partB_train_mat, self.ori_dir_partB_train_lab)
if len(os.listdir(self.ori_dir_partB_test_mat)) != len(os.listdir(self.ori_dir_partB_test_lab)):
self._convert_mat_to_json(self.ori_dir_partB_test_mat, self.ori_dir_partB_test_lab)
def _create_train_test(self, force_augmentation, kwargs):
slide_window_params = {
'displace': kwargs['displace'],
'size_x': kwargs['size_x'],
'size_y': kwargs['size_y'],
'people_thr': kwargs['people_thr']
}
noise_params = {'augment_noise': kwargs['augment_noise']}
light_params = {
'augment_light': kwargs['augment_light'],
'bright': kwargs['bright'],
'contrast': kwargs['contrast']
}
#shanghaiTech part A
self.augmented_dir_partA = osp.join(self.ori_dir_partA,
self.signature())
augment_data_A = False
if osp.exists(self.augmented_dir_partA):
print("'{}' already exists".format(self.augmented_dir_partA))
if force_augmentation:
augment_data_A = True
print("augmenting data anyway")
else:
augment_data_A = False
print("will not augmenting data")
else:
augment_data_A = True
os.makedirs(self.augmented_dir_partA)
aug_dir_partA_img = osp.join(self.augmented_dir_partA, "train_img")
aug_dir_partA_den = osp.join(self.augmented_dir_partA, "train_den")
aug_dir_partA_lab = osp.join(self.augmented_dir_partA, "train_lab")
test_dir_partA_img = osp.join(self.augmented_dir_partA, 'test_img')
test_dir_partA_den = osp.join(self.augmented_dir_partA, 'test_den')
test_dir_partA_lab = osp.join(self.augmented_dir_partA, 'test_lab')
mkdir_if_missing(aug_dir_partA_img)
mkdir_if_missing(aug_dir_partA_den)
mkdir_if_missing(aug_dir_partA_lab)
mkdir_if_missing(test_dir_partA_img)
mkdir_if_missing(test_dir_partA_den)
mkdir_if_missing(test_dir_partA_lab)
kwargs['name'] = 'shanghai-partA'
part_A_train_test = train_test_unit(aug_dir_partA_img,
aug_dir_partA_den,
test_dir_partA_img,
test_dir_partA_den, kwargs.copy())
self.train_test_set.append(part_A_train_test)
if augment_data_A:
ori_img_paths = [
osp.join(self.ori_dir_partA_train_img, file_name)
for file_name in sorted(
os.listdir(self.ori_dir_partA_train_img))
]
ori_lab_paths = [
osp.join(self.ori_dir_partA_train_lab, file_name)
for file_name in sorted(
os.listdir(self.ori_dir_partA_train_lab))
]
ori_den_paths = [
osp.join(self.ori_dir_partA_train_den, file_name)
for file_name in sorted(
os.listdir(self.ori_dir_partA_train_den))
]
augment(ori_img_paths, ori_lab_paths, ori_den_paths,
aug_dir_partA_img, aug_dir_partA_lab, aug_dir_partA_den,
slide_window_params, noise_params, light_params)
test_img_paths = [
osp.join(self.ori_dir_partA_test_img,
file_name) for file_name in sorted(
os.listdir(self.ori_dir_partA_test_img))
]
test_lab_paths = [
osp.join(self.ori_dir_partA_test_lab,
file_name) for file_name in sorted(
os.listdir(self.ori_dir_partA_test_lab))
]
test_den_paths = [
osp.join(self.ori_dir_partA_test_den,
file_name) for file_name in sorted(
os.listdir(self.ori_dir_partA_test_den))
]
copy_to_directory(test_img_paths, test_dir_partA_img)
copy_to_directory(test_lab_paths, test_dir_partA_lab)
copy_to_directory(test_den_paths, test_dir_partA_den)
#shanghaiTech part B
self.augmented_dir_partB = osp.join(self.ori_dir_partB,
self.signature())
augment_data_B = False
if osp.exists(self.augmented_dir_partB):
print("'{}' already exists".format(self.augmented_dir_partB))
if force_augmentation:
augment_data_B = True
print("augmenting data anyway")
else:
augment_data_B = False
print("will not augmenting data")
else:
augment_data_B = True
os.makedirs(self.augmented_dir_partB)
aug_dir_partB_img = osp.join(self.augmented_dir_partB, "train_img")
aug_dir_partB_den = osp.join(self.augmented_dir_partB, "train_den")
aug_dir_partB_lab = osp.join(self.augmented_dir_partB, "train_lab")
test_dir_partB_img = osp.join(self.augmented_dir_partB, 'test_img')
test_dir_partB_den = osp.join(self.augmented_dir_partB, 'test_den')
test_dir_partB_lab = osp.join(self.augmented_dir_partB, 'test_lab')
mkdir_if_missing(aug_dir_partB_img)
mkdir_if_missing(aug_dir_partB_den)
mkdir_if_missing(aug_dir_partB_lab)
mkdir_if_missing(test_dir_partB_img)
mkdir_if_missing(test_dir_partB_den)
mkdir_if_missing(test_dir_partB_lab)
kwargs['name'] = 'shanghai-partB'
part_B_train_test = train_test_unit(aug_dir_partB_img,
aug_dir_partB_den,
test_dir_partB_img,
test_dir_partB_den, kwargs.copy())
self.train_test_set.append(part_B_train_test)
if augment_data_B:
ori_img_paths = [
osp.join(self.ori_dir_partB_train_img, file_name)
for file_name in sorted(
os.listdir(self.ori_dir_partB_train_img))
]
ori_lab_paths = [
osp.join(self.ori_dir_partB_train_lab, file_name)
for file_name in sorted(
os.listdir(self.ori_dir_partB_train_lab))
]
ori_den_paths = [
osp.join(self.ori_dir_partB_train_den, file_name)
for file_name in sorted(
os.listdir(self.ori_dir_partB_train_den))
]
augment(ori_img_paths, ori_lab_paths, ori_den_paths,
aug_dir_partB_img, aug_dir_partB_lab, aug_dir_partB_den,
slide_window_params, noise_params, light_params)
test_img_paths = [
osp.join(self.ori_dir_partB_test_img,
file_name) for file_name in sorted(
os.listdir(self.ori_dir_partB_test_img))
]
test_lab_paths = [
osp.join(self.ori_dir_partB_test_lab,
file_name) for file_name in sorted(
os.listdir(self.ori_dir_partB_test_lab))
]
test_den_paths = [
osp.join(self.ori_dir_partB_test_den,
file_name) for file_name in sorted(
os.listdir(self.ori_dir_partB_test_den))
]
copy_to_directory(test_img_paths, test_dir_partB_img)
copy_to_directory(test_lab_paths, test_dir_partB_lab)
copy_to_directory(test_den_paths, test_dir_partB_den)
def train(train_test_unit, out_dir_root):
output_dir = osp.join(out_dir_root, train_test_unit.metadata['name'])
mkdir_if_missing(output_dir)
sys.stdout = Logger(osp.join(output_dir, 'log_train.txt'))
print("==========\nArgs:{}\n==========".format(args))
dataset_name = train_test_unit.metadata['name']
train_path = train_test_unit.train_dir_img
train_gt_path = train_test_unit.train_dir_den
val_path = train_test_unit.test_dir_img
val_gt_path = train_test_unit.test_dir_den
# training configuration
start_step = args.start_epoch
end_step = args.max_epoch
lr = args.lr
# log frequency
disp_interval = args.train_batch * 20
# ------------
rand_seed = args.seed
if rand_seed is not None:
np.random.seed(rand_seed)
torch.manual_seed(rand_seed)
torch.cuda.manual_seed(rand_seed)
# load net
net = CrowdCounter()
if not args.resume:
network.weights_normal_init(net, dev=0.01)
else:
# network.weights_normal_init(net, dev=0.01) #init all layers in case of partial net load
if args.resume[-3:] == '.h5':
pretrained_model = args.resume
else:
resume_dir = osp.join(args.resume, pu.metadata['name'])
pretrained_model = osp.join(resume_dir, 'best_model.h5')
network.load_net(pretrained_model, net)
print('Will apply fine tunning over', pretrained_model)
net.cuda()
net.train()
optimizer_d_large = torch.optim.Adam(filter(lambda p: p.requires_grad,
net.d_large.parameters()),
lr=lr,
betas=(args.beta1, args.beta2))
optimizer_d_small = torch.optim.Adam(filter(lambda p: p.requires_grad,
net.d_small.parameters()),
lr=lr,
betas=(args.beta1, args.beta2))
optimizer_g_large = torch.optim.Adam(filter(lambda p: p.requires_grad,
net.g_large.parameters()),
lr=lr,
betas=(args.beta1, args.beta2))
optimizer_g_small = torch.optim.Adam(filter(lambda p: p.requires_grad,
net.g_small.parameters()),
lr=lr,
betas=(args.beta1, args.beta2))
# training
train_loss = 0
step_cnt = 0
re_cnt = False
t = Timer()
t.tic()
# preprocess flags
overlap_test = True if args.overlap_test else False
data_loader = ImageDataLoader(train_path,
train_gt_path,
shuffle=True,
batch_size=args.train_batch,
test_loader=False)
data_loader_val = ImageDataLoader(val_path,
val_gt_path,
shuffle=False,
batch_size=1,
test_loader=True,
img_width=args.size_x,
img_height=args.size_y,
test_overlap=overlap_test)
best_mae = sys.maxsize
for epoch in range(start_step, end_step + 1):
step = 0
train_loss_gen_small = 0
train_loss_gen_large = 0
train_loss_dis_small = 0
train_loss_dis_large = 0
for blob in data_loader:
step = step + args.train_batch
im_data = blob['data']
gt_data = blob['gt_density']
idx_data = blob['idx']
im_data_norm = im_data / 127.5 - 1. # normalize between -1 and 1
gt_data = gt_data * args.den_factor
optimizer_d_large.zero_grad()
optimizer_d_small.zero_grad()
density_map = net(im_data_norm,
gt_data,
epoch=epoch,
mode="discriminator")
loss_d_small = net.loss_dis_small
loss_d_large = net.loss_dis_large
loss_d_small.backward()
loss_d_large.backward()
optimizer_d_small.step()
optimizer_d_large.step()
optimizer_g_large.zero_grad()
optimizer_g_small.zero_grad()
density_map = net(im_data_norm,
gt_data,
epoch=epoch,
mode="generator")
loss_g_small = net.loss_gen_small
loss_g_large = net.loss_gen_large
loss_g = net.loss_gen
loss_g.backward() # loss_g_large + loss_g_small
optimizer_g_small.step()
optimizer_g_large.step()
density_map /= args.den_factor
gt_data /= args.den_factor
train_loss_gen_small += loss_g_small.data.item()
train_loss_gen_large += loss_g_large.data.item()
train_loss_dis_small += loss_d_small.data.item()
train_loss_dis_large += loss_d_large.data.item()
step_cnt += 1
if step % disp_interval == 0:
duration = t.toc(average=False)
fps = step_cnt / duration
density_map = density_map.data.cpu().numpy()
train_batch_size = gt_data.shape[0]
gt_count = np.sum(gt_data.reshape(train_batch_size, -1),
axis=1)
et_count = np.sum(density_map.reshape(train_batch_size, -1),
axis=1)
if args.save_plots:
plot_save_dir = osp.join(output_dir, 'plot-results-train/')
mkdir_if_missing(plot_save_dir)
utils.save_results(im_data,
gt_data,
density_map,
idx_data,
plot_save_dir,
loss=args.loss)
print(
"epoch: {0}, step {1}/{5}, Time: {2:.4f}s, gt_cnt: {3:.4f}, et_cnt: {4:.4f}, mean_diff: {6:.4f}"
.format(epoch, step, 1. / fps, gt_count[0], et_count[0],
data_loader.num_samples,
np.mean(np.abs(gt_count - et_count))))
re_cnt = True
if re_cnt:
t.tic()
re_cnt = False
save_name = os.path.join(
output_dir, '{}_{}_{}.h5'.format(train_test_unit.to_string(),
dataset_name, epoch))
network.save_net(save_name, net)
# calculate error on the validation dataset
mae, mse = evaluate_model(save_name,
data_loader_val,
epoch=epoch,
den_factor=args.den_factor)
if mae < best_mae:
best_mae = mae
best_mse = mse
best_model = '{}_{}_{}.h5'.format(train_test_unit.to_string(),
dataset_name, epoch)
network.save_net(os.path.join(output_dir, "best_model.h5"), net)
print(
"Epoch: {0}, MAE: {1:.4f}, MSE: {2:.4f}, loss gen small: {3:.4f}, loss gen large: {4:.4f}, loss dis small: {5:.4f}, loss dis large: {6:.4f}, loss: {7:.4f}"
.format(
epoch, mae, mse, train_loss_gen_small, train_loss_gen_large,
train_loss_dis_small, train_loss_dis_large,
train_loss_gen_small + train_loss_gen_large +
train_loss_dis_small + train_loss_dis_large))
print("Best MAE: {0:.4f}, Best MSE: {1:.4f}, Best model: {2}".format(
best_mae, best_mse, best_model))
def evaluate_model(trained_model,
data_loader,
model,
save_test_results=False,
plot_save_dir="/tmp/",
den_scale_factor=1e3,
cam=False,
**kwargs):
mkdir_if_missing(plot_save_dir)
net = CrowdCounter(model=model, channel_param=kwargs['channel_param'])
network.load_net(trained_model, net)
net.cuda()
net.eval()
mae = 0.0
mse = 0.0
if cam:
net.net.cam = True
res = {
'img': [],
'et': [],
'gt': [],
'abs_err': [],
'err^2': [],
'err/gt': []
}
for blob in data_loader:
im_data = blob['data']
gt_data = blob['gt_density']
idx_data = blob['idx']
im_data_norm = im_data / 127.5 - 1.
density_map = net(im_data_norm)
if cam in ['grad-cam', 'all']:
apply_cam(net, blob, density_map,
osp.join(plot_save_dir, 'grad-cam'), 'grad-cam')
if cam in ['grad-cam++', 'all']:
apply_cam(net, blob, density_map,
osp.join(plot_save_dir, 'grad-cam++'), 'grad-cam++')
#if cam in ['eigen-cam', 'all']:
# apply_cam(net, blob, density_map, osp.join(plot_save_dir, 'eigen-cam'), 'eigen-cam')
if cam in ['layer-cam', 'all']:
apply_cam(net, blob, density_map,
osp.join(plot_save_dir, 'layer-cam'), 'layer-cam')
density_map = density_map.data.cpu().numpy()
density_map /= den_scale_factor
gt_count = np.sum(gt_data)
et_count = np.sum(density_map)
mae += abs(gt_count - et_count)
mse += ((gt_count - et_count) * (gt_count - et_count))
res['img'].append(idx_data)
res['et'].append(et_count)
res['gt'].append(gt_count)
res['abs_err'].append(abs(gt_count - et_count))
res['err^2'].append((gt_count - et_count) * (gt_count - et_count))
res['err/gt'].append(abs(gt_count - et_count) / gt_count)
if save_test_results:
print("Plotting results")
mkdir_if_missing(plot_save_dir)
save_results(im_data, gt_data, density_map, idx_data,
plot_save_dir)
mae = mae / data_loader.get_num_samples()
mse = np.sqrt(mse / data_loader.get_num_samples())
df = pd.DataFrame(res).sort_values('err/gt', axis=0, ascending=False)
df.to_csv(osp.join(plot_save_dir, 'error_by_img.csv'),
sep=',',
index=False,
float_format='%.5f')
return mae, mse
def train(train_test_unit, out_dir_root):
if args.model in ['mcnn4-gan', 'mcnn4-gan-skip', 'mcnn4-gan-u']:
train_gan(train_test_unit, out_dir_root, args)
return
output_dir = osp.join(out_dir_root, train_test_unit.metadata['name'])
mkdir_if_missing(output_dir)
output_dir_model = osp.join(output_dir, 'models')
mkdir_if_missing(output_dir_model)
if args.resume:
sys.stdout = Logger(osp.join(output_dir, 'log_train.txt'), mode='a')
else:
sys.stdout = Logger(osp.join(output_dir, 'log_train.txt'))
print("==========\nArgs:{}\n==========".format(args))
dataset_name = train_test_unit.metadata['name']
train_path = train_test_unit.train_dir_img
train_gt_path = train_test_unit.train_dir_den
val_path = train_test_unit.val_dir_img
val_gt_path = train_test_unit.val_dir_den
#training configuration
start_step = args.start_epoch
end_step = args.max_epoch
lr = args.lr
#log frequency
disp_interval = args.train_batch * 20
# ------------
rand_seed = args.seed
if rand_seed is not None:
np.random.seed(rand_seed)
torch.manual_seed(rand_seed)
torch.cuda.manual_seed(rand_seed)
best_mae = sys.maxsize # best mae
current_patience = 0
# load net
net = CrowdCounter(model=args.model)
if not args.resume:
network.weights_normal_init(net, dev=0.01)
else:
if args.resume[-3:] == '.h5': #don't use this option!
pretrained_model = args.resume
else:
resume_dir = osp.join(args.resume,
train_test_unit.metadata['name'])
if args.last_model:
pretrained_model = osp.join(resume_dir, 'last_model.h5')
f = open(osp.join(resume_dir, "current_values.bin"), "rb")
current_patience = pickle.load(f)
f.close()
else:
pretrained_model = osp.join(resume_dir, 'best_model.h5')
current_patience = 0
f = open(osp.join(resume_dir, "best_values.bin"), "rb")
best_mae, best_mse, best_model, _ = pickle.load(f)
f.close()
print(
"Best MAE: {0:.4f}, Best MSE: {1:.4f}, Best model: {2}, Current patience: {3}"
.format(best_mae, best_mse, best_model, current_patience))
network.load_net(pretrained_model, net)
print('Will apply fine tunning over', pretrained_model)
net.cuda()
net.train()
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
net.parameters()),
lr=lr)
# training
train_loss = 0
step_cnt = 0
re_cnt = False
t = Timer()
t.tic()
data_loader = ImageDataLoader(train_path,
train_gt_path,
shuffle=True,
batch_size=args.train_batch)
data_loader_val = ImageDataLoader(val_path,
val_gt_path,
shuffle=False,
batch_size=1)
for epoch in range(start_step, end_step + 1):
step = 0
train_loss = 0
for blob in data_loader:
optimizer.zero_grad()
step = step + args.train_batch
im_data = blob['data']
gt_data = blob['gt_density']
im_data_norm = im_data / 127.5 - 1. #normalize between -1 and 1
gt_data *= args.den_scale_factor
density_map = net(im_data_norm, gt_data=gt_data)
loss = net.loss
loss.backward()
optimizer.step()
train_loss += loss.data.item()
density_map = density_map.data.cpu().numpy()
density_map /= args.den_scale_factor
gt_data /= args.den_scale_factor
step_cnt += 1
if step % disp_interval == 0:
duration = t.toc(average=False)
fps = step_cnt / duration
train_batch_size = gt_data.shape[0]
gt_count = np.sum(gt_data.reshape(train_batch_size, -1),
axis=1)
et_count = np.sum(density_map.reshape(train_batch_size, -1),
axis=1)
print(
"epoch: {0}, step {1}/{5}, Time: {2:.4f}s, gt_cnt[0]: {3:.4f}, et_cnt[0]: {4:.4f}, mean_diff: {6:.4f}"
.format(epoch, step, 1. / fps, gt_count[0], et_count[0],
data_loader.num_samples,
np.mean(np.abs(gt_count - et_count))))
re_cnt = True
if re_cnt:
t.tic()
re_cnt = False
save_name = os.path.join(
output_dir_model, '{}_{}_{}.h5'.format(train_test_unit.to_string(),
dataset_name, epoch))
network.save_net(save_name, net)
network.save_net(os.path.join(output_dir, "last_model.h5"), net)
#calculate error on the validation dataset
mae, mse = evaluate_model(save_name,
data_loader_val,
model=args.model,
save_test_results=args.save_plots,
plot_save_dir=osp.join(
output_dir, 'plot-results-train/'),
den_scale_factor=args.den_scale_factor)
if mae < best_mae:
best_mae = mae
best_mse = mse
current_patience = 0
best_model = '{}_{}_{}.h5'.format(train_test_unit.to_string(),
dataset_name, epoch)
network.save_net(os.path.join(output_dir, "best_model.h5"), net)
f = open(os.path.join(output_dir, "best_values.bin"), "wb")
pickle.dump((best_mae, best_mse, best_model, current_patience), f)
f.close()
else:
current_patience += 1
f = open(os.path.join(output_dir, "current_values.bin"), "wb")
pickle.dump(current_patience, f)
f.close()
print("Epoch: {0}, MAE: {1:.4f}, MSE: {2:.4f}, loss: {3:.4f}".format(
epoch, mae, mse, train_loss))
print("Best MAE: {0:.4f}, Best MSE: {1:.4f}, Best model: {2}".format(
best_mae, best_mse, best_model))
print("Patience: {0}/{1}".format(current_patience, args.patience))
sys.stdout.close_open()
if current_patience > args.patience and args.patience > -1:
break
def _create_train_test(self, force_augmentation, **kwargs):
slide_window_params = {
'displace': kwargs['displace'],
'size_x': kwargs['size_x'],
'size_y': kwargs['size_y'],
'people_thr': kwargs['people_thr']
}
noise_params = {'augment_noise': kwargs['augment_noise']}
light_params = {
'augment_light': kwargs['augment_light'],
'bright': kwargs['bright'],
'contrast': kwargs['contrast']
}
self.augmented_dir = osp.join(self.root, self.signature())
# check if augmented data exists
augment_data = False
if osp.exists(self.augmented_dir):
print("'{}' already exists".format(self.augmented_dir))
if force_augmentation:
augment_data = True
print("augmenting data anyway, since force_augmentantion=True")
else:
augment_data = False
print("skipping augmentation, since force_augmentantion=False")
else:
augment_data = True
os.makedirs(self.augmented_dir)
#create new directories for augmentations
aug_dir_train_img = osp.join(self.augmented_dir, "train_img")
aug_dir_train_den = osp.join(self.augmented_dir, "train_den")
aug_dir_train_lab = osp.join(self.augmented_dir, "train_lab")
test_dir_img = osp.join(self.augmented_dir, "test_img")
test_dir_den = osp.join(self.augmented_dir, "test_den")
test_dir_lab = osp.join(self.augmented_dir, "test_lab")
val_dir_img = osp.join(self.augmented_dir, "val_img")
val_dir_den = osp.join(self.augmented_dir, "val_den")
val_dir_lab = osp.join(self.augmented_dir, "val_lab")
mkdir_if_missing(aug_dir_train_img)
mkdir_if_missing(aug_dir_train_den)
mkdir_if_missing(aug_dir_train_lab)
mkdir_if_missing(test_dir_img)
mkdir_if_missing(test_dir_den)
mkdir_if_missing(test_dir_lab)
mkdir_if_missing(val_dir_img)
mkdir_if_missing(val_dir_den)
mkdir_if_missing(val_dir_lab)
kwargs['name'] = 'jhu_crowd'
kwargs["val"] = True
kwargs["val_dir_img"] = val_dir_img
kwargs["val_dir_den"] = val_dir_den
train_test = train_test_unit(aug_dir_train_img, aug_dir_train_den,
test_dir_img, test_dir_den, kwargs.copy())
if augment_data and not (kwargs["skip_dir_check"]):
ori_img_paths = [
osp.join(self.ori_train_img, file_name)
for file_name in sorted(os.listdir(self.ori_train_img))
]
ori_lab_paths = [
osp.join(self.ori_train_lab, file_name)
for file_name in sorted(os.listdir(self.ori_train_lab))
]
ori_den_paths = [
osp.join(self.ori_train_den, file_name)
for file_name in sorted(os.listdir(self.ori_train_den))
]
augment(ori_img_paths, ori_lab_paths, ori_den_paths,
aug_dir_train_img, aug_dir_train_lab, aug_dir_train_den,
slide_window_params, noise_params, light_params)
test_img_paths = [
osp.join(self.ori_test_img, file_name)
for file_name in sorted(os.listdir(self.ori_test_img))
]
test_lab_paths = [
osp.join(self.ori_test_lab, file_name)
for file_name in sorted(os.listdir(self.ori_test_lab))
]
test_den_paths = [
osp.join(self.ori_test_den, file_name)
for file_name in sorted(os.listdir(self.ori_test_den))
]
copy_to_directory(test_img_paths, test_dir_img)
copy_to_directory(test_lab_paths, test_dir_lab)
copy_to_directory(test_den_paths, test_dir_den)
val_img_paths = [
osp.join(self.ori_val_img, file_name)
for file_name in sorted(os.listdir(self.ori_val_img))
]
val_lab_paths = [
osp.join(self.ori_val_lab, file_name)
for file_name in sorted(os.listdir(self.ori_val_lab))
]
val_den_paths = [
osp.join(self.ori_val_den, file_name)
for file_name in sorted(os.listdir(self.ori_val_den))
]
copy_to_directory(val_img_paths, val_dir_img)
copy_to_directory(val_lab_paths, val_dir_lab)
copy_to_directory(val_den_paths, val_dir_den)
self.train_test_set.append(train_test)
def train_gan(train_test_unit, out_dir_root, args):
output_dir = osp.join(out_dir_root, train_test_unit.metadata['name'])
mkdir_if_missing(output_dir)
output_dir_model = osp.join(output_dir, 'models')
mkdir_if_missing(output_dir_model)
if args.resume:
sys.stdout = Logger(osp.join(output_dir, 'log_train.txt'), mode='a')
plotter = LossPlotter(output_dir, mode='a')
else:
sys.stdout = Logger(osp.join(output_dir, 'log_train.txt'))
plotter = LossPlotter(output_dir, mode='w')
print("==========\nArgs:{}\n==========".format(args))
dataset_name = train_test_unit.metadata['name']
train_path = train_test_unit.train_dir_img
train_gt_path = train_test_unit.train_dir_den
val_path = train_test_unit.val_dir_img
val_gt_path = train_test_unit.val_dir_den
#training configuration
start_step = args.start_epoch
end_step = args.max_epoch
#log frequency
disp_interval = args.train_batch * 20
# ------------
rand_seed = args.seed
if rand_seed is not None:
np.random.seed(rand_seed)
torch.manual_seed(rand_seed)
torch.cuda.manual_seed(rand_seed)
best_mae = sys.maxsize # best mae
current_patience = 0
mse_criterion = nn.MSELoss()
# load net and optimizer
net = CrowdCounter(model=args.model, channel_param=args.channel_param)
net.cuda()
net.train()
#optimizerG = torch.optim.RMSprop(filter(lambda p: p.requires_grad, net.net.parameters()), lr=lr)
#optimizerD = torch.optim.RMSprop(filter(lambda p: p.requires_grad, net.gan_net.parameters()), lr=lrc)
optimizerG, optimizerD = get_optimizers(args, net)
if args.reduce_lr_on_plateau:
schedulerG = lr_scheduler.ReduceLROnPlateau(
optimizerG,
patience=args.scheduler_patience,
factor=args.scheduler_factor,
cooldown=args.scheduler_cooldown,
min_lr=args.min_lr,
verbose=True)
schedulerD = lr_scheduler.ReduceLROnPlateau(
optimizerD,
patience=args.scheduler_patience,
factor=args.scheduler_factor,
cooldown=args.scheduler_cooldown,
min_lr=args.min_lrc,
verbose=True)
elif args.step_lr:
schedulerG = lr_scheduler.StepLR(optimizerG,
args.scheduler_step_size,
gamma=args.scheduler_gamma,
verbose=True)
schedulerD = lr_scheduler.StepLR(optimizerD,
args.scheduler_step_size,
gamma=args.scheduler_gamma,
verbose=True)
if not args.resume:
network.weights_normal_init(net.net, dev=0.01)
else:
if args.resume[-3:] == '.h5': #don't use this option!
pretrained_model = args.resume
else:
resume_dir = osp.join(args.resume,
train_test_unit.metadata['name'])
if args.last_model:
pretrained_model = osp.join(resume_dir, 'last_model.h5')
f = open(osp.join(resume_dir, "current_values.bin"), "rb")
current_patience = pickle.load(f)
f.close()
f = torch.load(osp.join(resume_dir, 'optimizer.pth'))
optimizerD.load_state_dict(f['opt_d'])
optimizerG.load_state_dict(f['opt_g'])
if args.reduce_lr_on_plateau or args.step_lr:
schedulerD.load_state_dict(f['sch_d'])
schedulerG.load_state_dict(f['sch_g'])
else:
pretrained_model = osp.join(resume_dir, 'best_model.h5')
current_patience = 0
f = open(osp.join(resume_dir, "best_values.bin"), "rb")
best_mae, best_mse, best_model, _ = pickle.load(f)
f.close()
print(
"Best MAE: {0:.4f}, Best MSE: {1:.4f}, Best model: {2}, Current patience: {3}"
.format(best_mae, best_mse, best_model, current_patience))
network.load_net(pretrained_model, net)
print('Will apply fine tunning over', pretrained_model)
# training
train_lossG = 0
train_lossD = 0
step_cnt = 0
re_cnt = False
t = Timer()
t.tic()
# gan labels
real_label = 1
fake_label = 0
netD = net.gan_net
netG = net.net
data_loader = ImageDataLoader(train_path,
train_gt_path,
shuffle=True,
batch_size=args.train_batch,
den_scale=1)
data_loader_val = ImageDataLoader(val_path,
val_gt_path,
shuffle=False,
batch_size=1,
den_scale=1,
testing=True)
for epoch in range(start_step, end_step + 1):
step = 0
train_lossG = 0
train_lossD = 0
train_lossG_mse = 0
train_lossG_gan = 0
for blob in data_loader:
optimizerG.zero_grad()
optimizerD.zero_grad()
step = step + args.train_batch
im_data = blob['data']
gt_data = blob['gt_density']
im_data_norm_a = im_data / 127.5 - 1. #normalize between -1 and 1
gt_data_a = gt_data * args.den_scale_factor
errD_epoch = 0
for critic_epoch in range(args.ncritic):
im_data_norm = network.np_to_variable(im_data_norm_a,
is_cuda=True,
is_training=True)
gt_data = network.np_to_variable(gt_data_a,
is_cuda=True,
is_training=True)
netD.zero_grad()
netG.zero_grad()
#real data discriminator
b_size = gt_data.size(0)
output_real = netD(gt_data).view(-1)
#fake data discriminator
density_map = netG(im_data_norm)
output_fake = netD(density_map.detach()).view(-1)
errD = -(torch.mean(output_real) - torch.mean(output_fake))
errD.backward()
optimizerD.step()
for p in netD.parameters():
p.data.clamp_(-0.01, 0.01)
errD_epoch += errD.data.item()
errD_epoch /= args.ncritic
#Generator update
netG.zero_grad()
output_fake = netD(density_map).view(-1)
errG_gan = -torch.mean(output_fake)
errG_mse = mse_criterion(density_map, gt_data)
#errG = (1-args.alpha)*errG_mse + args.alpha*errG_gan
errG = errG_mse + args.alpha * errG_gan
errG.backward()
optimizerG.step()
train_lossG += errG.data.item()
train_lossG_mse += errG_mse.data.item()
train_lossG_gan += errG_gan.data.item()
train_lossD += errD_epoch
density_map = density_map.data.cpu().numpy()
density_map /= args.den_scale_factor
gt_data = gt_data.data.cpu().numpy()
gt_data /= args.den_scale_factor
step_cnt += 1
if step % disp_interval == 0:
duration = t.toc(average=False)
fps = step_cnt / duration
train_batch_size = gt_data.shape[0]
gt_count = np.sum(gt_data.reshape(train_batch_size, -1),
axis=1)
et_count = np.sum(density_map.reshape(train_batch_size, -1),
axis=1)
print(
"epoch: {0}, step {1}/{5}, Time: {2:.4f}s, gt_cnt[0]: {3:.4f}, et_cnt[0]: {4:.4f}, mean_diff: {6:.4f}"
.format(epoch, step, 1. / fps, gt_count[0], et_count[0],
data_loader.num_samples,
np.mean(np.abs(gt_count - et_count))))
re_cnt = True
if re_cnt:
t.tic()
re_cnt = False
#save model and optimizer
save_name = os.path.join(
output_dir_model, '{}_{}_{}.h5'.format(train_test_unit.to_string(),
dataset_name, epoch))
network.save_net(save_name, net)
network.save_net(os.path.join(output_dir, "last_model.h5"), net)
#calculate error on the validation dataset
mae, mse = evaluate_model(save_name,
data_loader_val,
model=args.model,
save_test_results=args.save_plots,
plot_save_dir=osp.join(
output_dir, 'plot-results-train/'),
den_scale_factor=args.den_scale_factor,
channel_param=args.channel_param)
if mae < best_mae:
best_mae = mae
best_mse = mse
current_patience = 0
best_model = '{}_{}_{}.h5'.format(train_test_unit.to_string(),
dataset_name, epoch)
network.save_net(os.path.join(output_dir, "best_model.h5"), net)
f = open(os.path.join(output_dir, "best_values.bin"), "wb")
pickle.dump((best_mae, best_mse, best_model, current_patience), f)
f.close()
else:
current_patience += 1
f = open(os.path.join(output_dir, "current_values.bin"), "wb")
pickle.dump(current_patience, f)
f.close()
# update lr
if args.reduce_lr_on_plateau:
schedulerD.step(train_lossG_mse)
schedulerG.step(train_lossG_mse)
elif args.step_lr:
schedulerD.step()
schedulerG.step()
optim_dict = {
"opt_d": optimizerD.state_dict(),
"opt_g": optimizerG.state_dict()
}
if args.reduce_lr_on_plateau or args.step_lr:
optim_dict['sch_d'] = schedulerD.state_dict()
optim_dict['sch_g'] = schedulerG.state_dict()
torch.save(optim_dict, os.path.join(output_dir, "optimizer.pth"))
plotter.report(train_lossG_mse, train_lossG_gan, train_lossD)
plotter.save()
plotter.plot()
print(
"Epoch: {0}, MAE: {1:.4f}, MSE: {2:.4f}, lossG: {3:.4f}, lossG_mse: {4:.4f}, lossG_gan: {5:.4f}, lossD: {6:.4f}"
.format(epoch, mae, mse, train_lossG, train_lossG_mse,
train_lossG_gan, train_lossD))
print("Best MAE: {0:.4f}, Best MSE: {1:.4f}, Best model: {2}".format(
best_mae, best_mse, best_model))
print("Patience: {0}/{1}".format(current_patience, args.patience))
sys.stdout.close_open()
if current_patience > args.patience and args.patience > -1:
break
def _create_labels(self):
mkdir_if_missing(self.ori_train_lab)
mkdir_if_missing(self.ori_train_den)
mkdir_if_missing(self.ori_val_lab)
mkdir_if_missing(self.ori_val_den)
mkdir_if_missing(self.ori_test_lab)
mkdir_if_missing(self.ori_test_den)
#convert csv to json
if len(os.listdir(self.ori_train_gt)) != len(
os.listdir(self.ori_train_lab)):
self._csv_to_json(self.ori_train_gt, self.ori_train_lab)
if len(os.listdir(self.ori_val_gt)) != len(os.listdir(
self.ori_val_lab)):
self._csv_to_json(self.ori_val_gt, self.ori_val_lab)
if len(os.listdir(self.ori_test_gt)) != len(
os.listdir(self.ori_test_lab)):
self._csv_to_json(self.ori_test_gt, self.ori_test_lab)
def _create_train_test(self, force_augmentation, kwargs):
slide_window_params = {
'displace': kwargs['displace'],
'size_x': kwargs['size_x'],
'size_y': kwargs['size_y'],
'people_thr': kwargs['people_thr']
}
noise_params = {'augment_noise': kwargs['augment_noise']}
light_params = {
'augment_light': kwargs['augment_light'],
'bright': kwargs['bright'],
'contrast': kwargs['contrast']
}
file_names = os.listdir(self.ori_dir_img)
file_names.sort()
img_names = []
img_ids = []
for file_name in file_names:
file_extention = file_name.split('.')[-1]
file_id = file_name[:len(file_name) - len(file_extention)]
if file_extention != 'png' and file_extention != 'jpg':
continue
img_names.append(file_name)
img_ids.append(file_id)
if len(img_names) != 50:
raise RuntimeError(
"UCF_CC_50 dataset expects 50 images, {} found".format(
len(img_names)))
self.augmented_dir = osp.join(self.root, self.signature())
augment_data = False
if osp.exists(self.augmented_dir):
print("'{}' already exists".format(self.augmented_dir))
if force_augmentation:
augment_data = True
print("augmenting data anyway")
else:
augment_data = False
print("will not augmenting data")
else:
augment_data = True
os.makedirs(self.augmented_dir)
#using 5 fold cross validation protocol
for fold in range(5):
fold_dir = osp.join(self.augmented_dir, 'fold{}'.format(fold + 1))
aug_train_dir_img = osp.join(fold_dir, 'train_img')
aug_train_dir_den = osp.join(fold_dir, 'train_den')
aug_train_dir_lab = osp.join(fold_dir, 'train_lab')
fold_test_dir_img = osp.join(fold_dir, 'test_img')
fold_test_dir_den = osp.join(fold_dir, 'test_den')
fold_test_dir_lab = osp.join(fold_dir, 'test_lab')
mkdir_if_missing(aug_train_dir_img)
mkdir_if_missing(aug_train_dir_den)
mkdir_if_missing(aug_train_dir_lab)
mkdir_if_missing(fold_test_dir_img)
mkdir_if_missing(fold_test_dir_den)
mkdir_if_missing(fold_test_dir_lab)
kwargs['name'] = 'ucf-fold{}'.format(fold + 1)
train_test = train_test_unit(aug_train_dir_img, aug_train_dir_den,
fold_test_dir_img, fold_test_dir_den,
kwargs.copy())
self.train_test_set.append(train_test)
if augment_data:
test_img = img_names[fold * 10:(fold + 1) * 10]
test_ids = img_ids[fold * 10:(fold + 1) * 10]
test_den_paths = [
osp.join(self.ori_dir_den, img_id + 'npy')
for img_id in test_ids
]
test_lab_paths = [
osp.join(self.ori_dir_lab, img_id + 'json')
for img_id in test_ids
]
test_img_paths = [
osp.join(self.ori_dir_img, img) for img in test_img
]
train_img = sorted(list(set(img_names) - set(test_img)))
train_ids = sorted(list(set(img_ids) - set(test_ids)))
train_den_paths = [
osp.join(self.ori_dir_den, img_id + 'npy')
for img_id in train_ids
]
train_lab_paths = [
osp.join(self.ori_dir_lab, img_id + 'json')
for img_id in train_ids
]
train_img_paths = [
osp.join(self.ori_dir_img, img) for img in train_img
]
#augment train data
print("Augmenting {}".format(kwargs['name']))
augment(train_img_paths, train_lab_paths, train_den_paths,
aug_train_dir_img, aug_train_dir_lab,
aug_train_dir_den, slide_window_params, noise_params,
light_params)
copy_to_directory(test_den_paths, fold_test_dir_den)
copy_to_directory(test_lab_paths, fold_test_dir_lab)
copy_to_directory(test_img_paths, fold_test_dir_img)
def train(train_test_unit, out_dir_root):
output_dir = osp.join(out_dir_root, train_test_unit.metadata['name'])
mkdir_if_missing(output_dir)
sys.stdout = Logger(osp.join(output_dir, 'log_train.txt'))
print("==========\nArgs:{}\n==========".format(args))
dataset_name = train_test_unit.metadata['name']
train_path = train_test_unit.train_dir_img
train_gt_path = train_test_unit.train_dir_den
val_path = train_test_unit.test_dir_img
val_gt_path = train_test_unit.test_dir_den
#training configuration
start_step = args.start_epoch
end_step = args.max_epoch
lr = args.lr
#log frequency
disp_interval = args.train_batch * 20
# ------------
rand_seed = args.seed
if rand_seed is not None:
np.random.seed(rand_seed)
torch.manual_seed(rand_seed)
torch.cuda.manual_seed(rand_seed)
# load net
net = CrowdCounter()
if not args.resume:
network.weights_normal_init(net, dev=0.01)
else:
if args.resume[-3:] == '.h5':
pretrained_model = args.resume
else:
resume_dir = osp.join(args.resume,
train_test_unit.metadata['name'])
pretrained_model = osp.join(resume_dir, 'best_model.h5')
network.load_net(pretrained_model, net)
print('Will apply fine tunning over', pretrained_model)
net.cuda()
net.train()
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
net.parameters()),
lr=lr)
# training
train_loss = 0
step_cnt = 0
re_cnt = False
t = Timer()
t.tic()
data_loader = ImageDataLoader(train_path,
train_gt_path,
shuffle=True,
batch_size=args.train_batch)
data_loader_val = ImageDataLoader(val_path,
val_gt_path,
shuffle=False,
batch_size=1)
best_mae = sys.maxsize
for epoch in range(start_step, end_step + 1):
step = 0
train_loss = 0
for blob in data_loader:
optimizer.zero_grad()
step = step + args.train_batch
im_data = blob['data']
gt_data = blob['gt_density']
im_data_norm = im_data / 255.0
density_map = net(im_data_norm, gt_data=gt_data)
loss = net.loss
loss.backward()
optimizer.step()
train_loss += loss.data.item()
density_map = density_map.data.cpu().numpy()
step_cnt += 1
if step % disp_interval == 0:
duration = t.toc(average=False)
fps = step_cnt / duration
train_batch_size = gt_data.shape[0]
gt_count = np.sum(gt_data.reshape(train_batch_size, -1),
axis=1)
et_count = np.sum(density_map.reshape(train_batch_size, -1),
axis=1)
print(
"epoch: {0}, step {1}/{5}, Time: {2:.4f}s, gt_cnt[0]: {3:.4f}, et_cnt[0]: {4:.4f}, mean_diff: {6:.4f}"
.format(epoch, step, 1. / fps, gt_count[0], et_count[0],
data_loader.num_samples,
np.mean(np.abs(gt_count - et_count))))
re_cnt = True
if re_cnt:
t.tic()
re_cnt = False
save_name = os.path.join(
output_dir, '{}_{}_{}.h5'.format(train_test_unit.to_string(),
dataset_name, epoch))
network.save_net(save_name, net)
#calculate error on the validation dataset
mae, mse = evaluate_model(save_name,
data_loader_val,
save_test_results=args.save_plots,
plot_save_dir=osp.join(
output_dir, 'plot-results-train/'))
if mae < best_mae:
best_mae = mae
best_mse = mse
best_model = '{}_{}_{}.h5'.format(train_test_unit.to_string(),
dataset_name, epoch)
network.save_net(os.path.join(output_dir, "best_model.h5"), net)
print("Epoch: {0}, MAE: {1:.4f}, MSE: {2:.4f}, loss: {3:.4f}".format(
epoch, mae, mse, train_loss))
print("Best MAE: {0:.4f}, Best MSE: {1:.4f}, Best model: {2}".format(
best_mae, best_mse, best_model))
def _create_train_test(self, force_augmentation, kwargs):
slide_window_params = {
'displace': kwargs['displace'],
'size_x': kwargs['size_x'],
'size_y': kwargs['size_y'],
'people_thr': kwargs['people_thr']
}
noise_params = {'augment_noise': kwargs['augment_noise']}
light_params = {
'augment_light': kwargs['augment_light'],
'bright': kwargs['bright'],
'contrast': kwargs['contrast']
}
downscale = kwargs['downscale']
ten_crop = kwargs['ten_crop_mode']
if downscale and ten_crop:
raise Exception(
'five crop and dowscale mode cannot be used at the same time, choose one.'
)
#shanghaiTech part A
self.augmented_dir_partA = osp.join(self.ori_dir_partA,
self.signature())
augment_data_A = False
if osp.exists(self.augmented_dir_partA):
print("'{}' already exists".format(self.augmented_dir_partA))
if force_augmentation:
augment_data_A = True
print("augmenting data anyway")
else:
augment_data_A = False
print("will not augmenting data")
else:
augment_data_A = True
os.makedirs(self.augmented_dir_partA)
aug_dir_partA_img = osp.join(self.augmented_dir_partA, "train_img")
aug_dir_partA_den = osp.join(self.augmented_dir_partA, "train_den")
aug_dir_partA_lab = osp.join(self.augmented_dir_partA, "train_lab")
test_dir_partA_img = osp.join(self.augmented_dir_partA, 'test_img')
test_dir_partA_den = osp.join(self.augmented_dir_partA, 'test_den')
test_dir_partA_lab = osp.join(self.augmented_dir_partA, 'test_lab')
mkdir_if_missing(aug_dir_partA_img)
mkdir_if_missing(aug_dir_partA_den)
mkdir_if_missing(aug_dir_partA_lab)
mkdir_if_missing(test_dir_partA_img)
mkdir_if_missing(test_dir_partA_den)
mkdir_if_missing(test_dir_partA_lab)
kwargs['name'] = 'shanghai-partA'
kwargs["val"] = False
part_A_train_test = train_test_unit(aug_dir_partA_img,
aug_dir_partA_den,
test_dir_partA_img,
test_dir_partA_den, kwargs.copy())
self.train_test_set.append(part_A_train_test)
if augment_data_A:
ori_img_paths = [
osp.join(self.ori_dir_partA_train_img, file_name)
for file_name in sorted(
os.listdir(self.ori_dir_partA_train_img))
]
ori_lab_paths = [
osp.join(self.ori_dir_partA_train_lab, file_name)
for file_name in sorted(
os.listdir(self.ori_dir_partA_train_lab))
]
ori_den_paths = [
osp.join(self.ori_dir_partA_train_den, file_name)
for file_name in sorted(
os.listdir(self.ori_dir_partA_train_den))
]
augment(ori_img_paths,
ori_lab_paths,
ori_den_paths,
aug_dir_partA_img,
aug_dir_partA_lab,
aug_dir_partA_den,
slide_window_params,
noise_params,
light_params,
downscale_img=downscale,
ten_crop=ten_crop)
test_img_paths = [
osp.join(self.ori_dir_partA_test_img,
file_name) for file_name in sorted(
os.listdir(self.ori_dir_partA_test_img))
]
test_lab_paths = [
osp.join(self.ori_dir_partA_test_lab,
file_name) for file_name in sorted(
os.listdir(self.ori_dir_partA_test_lab))
]
test_den_paths = [
osp.join(self.ori_dir_partA_test_den,
file_name) for file_name in sorted(
os.listdir(self.ori_dir_partA_test_den))
]
if downscale:
downscale_test_images(test_img_paths, test_dir_partA_img,
test_lab_paths, test_dir_partA_lab,
test_den_paths, test_dir_partA_den,
slide_window_params['size_x'],
slide_window_params['size_y'])
elif ten_crop:
resize_test_img_ten_crop_mode(
test_img_paths, test_dir_partA_img, test_lab_paths,
test_dir_partA_lab, test_den_paths, test_dir_partA_den)
else:
copy_to_directory(test_img_paths, test_dir_partA_img)
copy_to_directory(test_lab_paths, test_dir_partA_lab)
copy_to_directory(test_den_paths, test_dir_partA_den)
#shanghaiTech part B
self.augmented_dir_partB = osp.join(self.ori_dir_partB,
self.signature())
augment_data_B = False
if osp.exists(self.augmented_dir_partB):
print("'{}' already exists".format(self.augmented_dir_partB))
if force_augmentation:
augment_data_B = True
print("augmenting data anyway")
else:
augment_data_B = False
print("will not augmenting data")
else:
augment_data_B = True
os.makedirs(self.augmented_dir_partB)
aug_dir_partB_img = osp.join(self.augmented_dir_partB, "train_img")
aug_dir_partB_den = osp.join(self.augmented_dir_partB, "train_den")
aug_dir_partB_lab = osp.join(self.augmented_dir_partB, "train_lab")
test_dir_partB_img = osp.join(self.augmented_dir_partB, 'test_img')
test_dir_partB_den = osp.join(self.augmented_dir_partB, 'test_den')
test_dir_partB_lab = osp.join(self.augmented_dir_partB, 'test_lab')
mkdir_if_missing(aug_dir_partB_img)
mkdir_if_missing(aug_dir_partB_den)
mkdir_if_missing(aug_dir_partB_lab)
mkdir_if_missing(test_dir_partB_img)
mkdir_if_missing(test_dir_partB_den)
mkdir_if_missing(test_dir_partB_lab)
kwargs['name'] = 'shanghai-partB'
kwargs["val"] = False
part_B_train_test = train_test_unit(aug_dir_partB_img,
aug_dir_partB_den,
test_dir_partB_img,
test_dir_partB_den, kwargs.copy())
self.train_test_set.append(part_B_train_test)
if augment_data_B:
ori_img_paths = [
osp.join(self.ori_dir_partB_train_img, file_name)
for file_name in sorted(
os.listdir(self.ori_dir_partB_train_img))
]
ori_lab_paths = [
osp.join(self.ori_dir_partB_train_lab, file_name)
for file_name in sorted(
os.listdir(self.ori_dir_partB_train_lab))
]
ori_den_paths = [
osp.join(self.ori_dir_partB_train_den, file_name)
for file_name in sorted(
os.listdir(self.ori_dir_partB_train_den))
]
augment(ori_img_paths,
ori_lab_paths,
ori_den_paths,
aug_dir_partB_img,
aug_dir_partB_lab,
aug_dir_partB_den,
slide_window_params,
noise_params,
light_params,
downscale_img=downscale,
ten_crop=ten_crop)
test_img_paths = [
osp.join(self.ori_dir_partB_test_img,
file_name) for file_name in sorted(
os.listdir(self.ori_dir_partB_test_img))
]
test_lab_paths = [
osp.join(self.ori_dir_partB_test_lab,
file_name) for file_name in sorted(
os.listdir(self.ori_dir_partB_test_lab))
]
test_den_paths = [
osp.join(self.ori_dir_partB_test_den,
file_name) for file_name in sorted(
os.listdir(self.ori_dir_partB_test_den))
]
if downscale:
downscale_test_images(test_img_paths, test_dir_partB_img,
test_lab_paths, test_dir_partB_lab,
test_den_paths, test_dir_partB_den,
slide_window_params['size_x'],
slide_window_params['size_y'])
elif ten_crop:
resize_test_img_ten_crop_mode(
test_img_paths, test_dir_partB_img, test_lab_paths,
test_dir_partB_lab, test_den_paths, test_dir_partB_den)
else:
copy_to_directory(test_img_paths, test_dir_partB_img)
copy_to_directory(test_lab_paths, test_dir_partB_lab)
copy_to_directory(test_den_paths, test_dir_partB_den)
