def validate(val_loader, epoch, model, criterion, args, save_path=False): # keeping statistics batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() acces = AverageMeter() mse = AverageMeter() mse_hip = AverageMeter() mse_coxa = AverageMeter() mse_femur = AverageMeter() mse_tibia = AverageMeter() mse_tarsus = AverageMeter() mse_jump = AverageMeter() avg_local_max = AverageMeter() # predictions and score maps num_cameras = 7 predictions = np.zeros( shape=( num_cameras + 1, val_loader.dataset.greatest_image_id() + 1, config["num_predict"], 2, ), dtype=np.float32, ) # num_cameras+1 for the mirrored camera 3 getLogger('df3d').debug("Predictions shape {}".format(predictions.shape)) if save_path is not None: print("SAVE_PATH:") print(save_path) unlabeled_folder_replace = save_path.replace("/", "-") # score_map_filename = os.path.join(args.data_folder, "./heatmap_{}.pkl".format(unlabeled_folder_replace)) score_map_filename = os.path.join( args.data_folder, "{}".format(save_path), args.output_folder, "heatmap_{}.pkl".format(unlabeled_folder_replace), ) score_map_path = Path(score_map_filename) score_map_path.parent.mkdir(exist_ok=True, parents=True) score_map_arr = np.memmap( score_map_filename, dtype="float32", mode="w+", shape=( num_cameras + 1, val_loader.dataset.greatest_image_id() + 1, config["num_predict"], args.hm_res[0], args.hm_res[1], ), ) # num_cameras+1 for the mirrored camera 3 # switch to evaluate mode model.eval() end = time.time() bar = Bar("Processing", max=len(val_loader)) #if logging.getLogger('df3d').isEnabledFor(logging.INFO) else NoOutputBar() bar.start() for i, (inputs, target, meta) in enumerate(val_loader): # measure data loading time data_time.update(time.time() - end) target = target.cuda(non_blocking=True) input_var = torch.autograd.Variable(inputs.cuda()) target_var = torch.autograd.Variable(target) # compute output with torch.no_grad(): output = model(input_var) score_map = output[-1].data.cpu() # loss = criterion(output[-1], target_var) acc = accuracy(score_map.cpu(), target.cpu(), args.acc_joints) # generate predictions if save_path is not None: for n, cam_id, cam_read_id, img_id in zip( range(score_map.size(0)), meta["cid"], meta["cam_read_id"], meta["pid"] ): smap = to_numpy(score_map[n, :, :, :]) score_map_arr[cam_read_id, img_id, :] = smap pr = Camera.hm_to_pred( score_map_arr[cam_read_id, img_id, :], threshold_abs=0.0 ) predictions[cam_read_id, img_id, :] = pr # measure accuracy and record loss mse_err = mse_acc(target_var.data.cpu(), score_map) mse.update(torch.mean(mse_err[args.acc_joints, :]), inputs.size(0)) mse_hip.update( torch.mean(mse_err[np.arange(0, mse_err.shape[0], 5), :]), inputs.size(0) ) mse_coxa.update( torch.mean(mse_err[np.arange(1, mse_err.shape[0], 5), :]), inputs.size(0) ) mse_femur.update( torch.mean(mse_err[np.arange(2, mse_err.shape[0], 5), :]), inputs.size(0) ) mse_tibia.update( torch.mean(mse_err[np.arange(3, mse_err.shape[0], 5), :]), inputs.size(0) ) mse_tarsus.update( torch.mean(mse_err[np.arange(4, mse_err.shape[0], 5), :]), inputs.size(0) ) losses.update(0, inputs.size(0)) acces.update(acc[0], inputs.size(0)) jump_thr = args.hm_res[0] * 5.0 / 64.0 mse_jump.update( np.mean(np.ravel(mse_err[args.acc_joints, :] > jump_thr)) * 100.0, inputs.size(0), ) if args.unlabeled and i < args.num_output_image: drosophila_img = drosophila_image_overlay( inputs, score_map, args.hm_res, 3, np.arange(config["num_predict"]), img_id=int(meta["image_name"][0].split("_")[-1]), ) template = "{}_overlay_val_epoch_{}_{}_{:06d}_{}.jpg" folder_name = meta["folder_name"][0] folder_name = ( folder_name.replace(" ", "_").replace("\\", "-").replace("/", "-") ) save_image_name = template.format( "left" if meta["cid"][0] == 0 else "right", epoch, folder_name, meta["pid"][0], meta["cid"][0], ) save_image_name = save_image_name.replace(" ", "_ ").replace("\\", "-") save_image( img_path=os.path.join(args.checkpoint_image_dir, save_image_name), img=drosophila_img, ) elif i < args.num_output_image: drosophila_img = drosophila_image_overlay( inputs, score_map, args.hm_res, 3, np.arange(config["num_predict"]) ) folder_name = meta["folder_name"][0] template = "./overlay_val_epoch_{}_{}_{:06d}_{}.jpg" save_image_name = template.format( epoch, folder_name.replace("/", "-"), meta["pid"][0], meta["cid"][0] ) save_image_name = save_image_name.replace(" ", "_ ").replace("\\", "-") save_image( img_path=os.path.join(args.checkpoint_image_dir, save_image_name), img=drosophila_img, ) # measure elapsed time batch_time.update(time.time() - end) end = time.time() # plot progress bar.suffix = "({batch}/{size}) D:{data:.6f}s|B:{bt:.3f}s|L:{loss:.8f}|Acc:{acc: .4f}|Mse:{mse: .3f}|F-T:{mse_femur: .3f}|T-Tar:{mse_tibia: .3f}|Tar-tip:{mse_tarsus: .3f}".format( batch=i + 1, size=len(val_loader), data=data_time.val, bt=batch_time.val, total=bar.elapsed_td, eta=bar.eta_td, loss=losses.avg, acc=acces.avg, mse=mse.avg, mse_hip=mse_hip.avg, mse_femur=mse_femur.avg, mse_tibia=mse_tibia.avg, mse_tarsus=mse_tarsus.avg, mse_coxa=mse_coxa.avg, mse_jump=mse_jump.avg, local_max_count=avg_local_max.avg, ) bar.next() bar.finish() if save_path is not None: score_map_arr.flush() else: score_map_arr = None return losses.avg, acces.avg, predictions, score_map_arr, mse.avg, mse_jump.avg
def step(loader, model, optimizer, mode, heatmap, epoch, num_classes, acc_joints): keys_am = [ "batch_time", "data_time", "losses", "acces", "mse", "body_coxa", "coxa_femur", "femur_tibia", "tibia_tarsus", "tarsus_tip", ] am = {k: AverageMeter() for k in keys_am} if mode == mode.train: model.train() elif mode == mode.test: model.eval() start = time.time() bar = (Bar("Processing", max=len(loader)) if logger.debug_enabled() else NoOutputBar()) bar.start() predictions = np.zeros( shape=( config["num_cameras"] + 1, loader.dataset.greatest_image_id() + 1, config["num_predict"], 2, ), dtype=np.float32, ) # num_cameras+1 for the mirrored camera 3 for i, (inputs, target, meta) in enumerate(loader): np.random.seed() am["data_time"].update(time.time() - start) input_var = torch.autograd.Variable(on_cuda(inputs)) target_var = torch.autograd.Variable(on_cuda(target, non_blocking=True)) output = model(input_var) heatmap_batch = output[-1].data.cpu() for n, cam_read_id, img_id in zip(range(heatmap_batch.size(0)), meta["cam_read_id"], meta["pid"]): smap = to_numpy(heatmap_batch[n, :, :, :]) pr = Camera.hm_to_pred(smap, threshold_abs=0.0) predictions[cam_read_id, img_id, :] = pr if heatmap is not None: heatmap[cam_read_id, img_id, :] = smap # loss = df3dLoss(output, target_var, meta["joint_exists"], num_classes) if mode == mode.train: optimizer.zero_grad() loss.backward() optimizer.step() acc = accuracy(heatmap_batch, target, acc_joints) # am["acces"].add(acc) mse_err = mse_acc(target_var.data.cpu(), heatmap_batch) # am["mse"].add(mse_err) am["batch_time"].update(time.time() - start) start = time.time() bar.suffix = "{epoch} | ({batch}/{size}) D:{data:.6f}s|B:{bt:.3f}s|L:{loss:.8f}|Acc:{acc: .4f}|Mse:{mse: .3f} F-T:{mse_femur: .3f} T-Tar:{mse_tibia: .3f} Tar-tip:{mse_tarsus: .3f}".format( epoch=epoch, batch=i + 1, size=len(loader), data=am["data_time"].val, bt=am["batch_time"].val, total=bar.elapsed_td, eta=bar.eta_td, loss=am["losses"].avg, acc=am["acces"].avg, mse=am["mse"].avg, mse_hip=am["body_coxa"].avg, mse_femur=am["coxa_femur"].avg, mse_tibia=am["femur_tibia"].avg, mse_tarsus=am["tarsus_tip"].avg, ) bar.next() bar.finish() return (predictions, heatmap, am["losses"].avg, am["acces"].avg, am["mse"].avg)
def train(train_loader, epoch, model, optimizer, criterion, args): batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() acces = AverageMeter() mse = AverageMeter() mse_hip = AverageMeter() mse_coxa = AverageMeter() mse_femur = AverageMeter() mse_tibia = AverageMeter() mse_tarsus = AverageMeter() mse_jump = AverageMeter() avg_local_max = AverageMeter() # to later save the disk predictions = dict() # switch to train mode model.train() end = time.time() bar = Bar("Processing", max=len(train_loader))# if logging.getLogger('df3d').isEnabledFor(logging.DEBUG) else NoOutputBar() bar.start() for i, (inputs, target, meta) in enumerate(train_loader): # reset seed for truly random transformations on input np.random.seed() # measure data loading time data_time.update(time.time() - end) input_var = ( torch.autograd.Variable(inputs.cuda()) if torch.cuda.is_available() else torch.autograd.Variable(inputs) ) target_var = torch.autograd.Variable(target.cuda(non_blocking=True)) # compute output output = model(input_var) score_map = output[-1].data.cpu() loss_weight = torch.ones((inputs.size(0), args.num_classes, 1, 1)) if np.any(np.array(meta["joint_exists"]) == 0): batch_index = (np.logical_not(meta["joint_exists"])).nonzero()[:, 0] joint_index = (np.logical_not(meta["joint_exists"])).nonzero()[:, 1] loss_weight[batch_index, joint_index, :, :] = 0.0 # logger.debug(loss_weight) loss_weight.requires_grad = True loss_weight = loss_weight.cuda() loss = weighted_mse_loss(output[0], target_var, weights=loss_weight) for j in range(1, len(output)): loss += weighted_mse_loss(output[j], target_var, weights=loss_weight) # Compute gradient and do SGD step optimizer.zero_grad() loss.backward() optimizer.step() acc = accuracy(score_map, target, args.acc_joints) # measure accuracy and record loss mse_err = mse_acc(target_var.data.cpu(), score_map) getLogger('df3d').debug(f'{mse_err.size()}, {mse_err[:, 0] > 50}, {meta["image_name"][0]}') mse.update(torch.mean(mse_err[args.acc_joints, :]), inputs.size(0)) # per joint mse_hip.update(torch.mean(mse_err[np.arange(0, 15, 5), :]), inputs.size(0)) mse_coxa.update(torch.mean(mse_err[np.arange(1, 15, 5), :]), inputs.size(0)) mse_femur.update(torch.mean(mse_err[np.arange(2, 15, 5), :]), inputs.size(0)) mse_tibia.update(torch.mean(mse_err[np.arange(3, 15, 5), :]), inputs.size(0)) mse_tarsus.update(torch.mean(mse_err[np.arange(4, 15, 5), :]), inputs.size(0)) losses.update(loss.item(), inputs.size(0)) acces.update(acc[0], inputs.size(0)) jump_thr = args.hm_res[0] * 5.0 / 64.0 mse_jump.update( np.mean(np.ravel(mse_err[args.acc_joints, :] > jump_thr)) * 100.0, inputs.size(0), ) if i < args.num_output_image: drosophila_img = drosophila_image_overlay( inputs, target, args.hm_res, 3, args.train_joints ) folder_name = ( meta["folder_name"][0].replace("/", "-") + meta["folder_name"][1].replace("/", "-") + meta["folder_name"][2].replace("/", "-") ) image_name = meta["image_name"][0] template = "./train_gt_{}_{}_{:06d}_{}_joint_{}.jpg" save_image( img_path=os.path.join( args.checkpoint_image_dir, template.format( epoch, folder_name, meta["pid"][0], meta["cid"][0], [] ), ), img=drosophila_img, ) if i < args.num_output_image and False: drosophila_img = drosophila_image_overlay( inputs, score_map, args.hm_res, 3, args.train_joints, img_id=int(meta["image_name"][0].split("_")[-1]), ) template = "{}_train_pred_epoch_{}_{}_{:06d}_{}.jpg" folder_name = meta["folder_name"][0] folder_name = ( folder_name.replace(" ", "_").replace("\\", "-").replace("/", "-") ) save_image_name = template.format( "left" if meta["cid"][0] == 0 else "right", epoch, folder_name, meta["pid"][0], meta["cid"][0], ) save_image_name = save_image_name.replace(" ", "_ ").replace("\\", "-") save_image( img_path=os.path.join(args.checkpoint_image_dir, save_image_name), img=drosophila_img, ) # measure elapsed time batch_time.update(time.time() - end) end = time.time() # plot progress bar.suffix = "({batch}/{size}) D:{data:.6f}s|B:{bt:.3f}s|L:{loss:.8f}|Acc:{acc: .4f}|Mse:{mse: .3f} F-T:{mse_femur: .3f} T-Tar:{mse_tibia: .3f} Tar-tip:{mse_tarsus: .3f} Jump:{mse_jump: .4f}".format( batch=i + 1, size=len(train_loader), data=data_time.val, bt=batch_time.val, total=bar.elapsed_td, eta=bar.eta_td, loss=losses.avg, acc=acces.avg, mse=mse.avg, mse_hip=mse_hip.avg, mse_femur=mse_femur.avg, mse_tibia=mse_tibia.avg, mse_tarsus=mse_tarsus.avg, mse_coxa=mse_coxa.avg, mse_jump=mse_jump.avg, ) bar.next() bar.finish() return losses.avg, acces.avg, predictions, mse.avg, mse_jump.avg