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 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 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