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)