def main(): device = torch.device('cpu') parser = argparse.ArgumentParser(description="ReID Baseline Training") parser.add_argument( "--config_file", default="./configs/debug.yml", help="path to config file", type=str ) parser.add_argument("opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER) args = parser.parse_args() num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1 if args.config_file != "": cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() ## build network # model = Baseline(num_classes, cfg.MODEL.LAST_STRIDE, cfg.MODEL.PRETRAIN_PATH, cfg.MODEL.NECK, # cfg.TEST.NECK_FEAT, cfg.MODEL.NAME, cfg.MODEL.PRETRAIN_CHOICE, # cfg) num_classes = 6*5 # num cameras = 6 * 5 bg for camera model = build_model(cfg, num_classes)
def main(): parser = argparse.ArgumentParser(description="Dense Correspondence") parser.add_argument( "--config-file", default="", metavar="FILE", help="Path to config file", type=str, ) parser.add_argument( "--skip-test", dest="skip_test", help="Do not test the final model", action="store_true", ) parser.add_argument( '--test_model', dest='test_model', action='store_true' ) parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER ) args = parser.parse_args() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() output_dir = os.path.join(cfg.MODEL_DIR, cfg.MODEL.NAME) if output_dir: os.makedirs(output_dir, exist_ok=True) if cfg.MODEL.TEST: # build model model = build_matching_model(cfg) model.to(cfg.MODEL.DEVICE) model = torch.nn.DataParallel(model) # load pretrained parameters model_dir = os.path.join(cfg.MODEL_DIR, cfg.MODEL.NAME) checkpointer = Checkpointer(cfg, model, save_dir=model_dir) _ = checkpointer.load(cfg.MODEL.WEIGHT) test(cfg, model) else: model = train(cfg) if not args.skip_test: test(cfg, model)
def main(): parser = argparse.ArgumentParser(description="ReID Baseline Inference") parser.add_argument( "--config_file", default="/media/bi/Data/Mycode/car_demo/ai_city/configs/submit.yml", help="path to config file", type=str) parser.add_argument("opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER) args = parser.parse_args() num_gpus = int( os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1 if args.config_file != "": cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() output_dir = cfg.OUTPUT_DIR if output_dir and not os.path.exists(output_dir): mkdir(output_dir) logger = setup_logger("reid_baseline", output_dir, 0) logger.info("Using {} GPUS".format(num_gpus)) logger.info(args) if args.config_file != "": logger.info("Loaded configuration file {}".format(args.config_file)) with open(args.config_file, 'r') as cf: config_str = "\n" + cf.read() logger.info(config_str) logger.info("Running with config:\n{}".format(cfg)) if cfg.MODEL.DEVICE == "cuda": os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID cudnn.benchmark = True train_loader, val_loader, num_query, num_classes, dataset = make_data_loader( cfg) model = build_model(cfg, num_classes) model.load_param(cfg.TEST.WEIGHT) ## read meta information dataset = VeRi(cfg.DATASETS.ROOT_DIR) print("test_tracks", dataset.test_tracks) indices_np = inference(cfg, model, val_loader, num_query, dataset) #write_result(indices_np, os.path.dirname(cfg.TEST.WEIGHT), topk=100) write_result_with_track(indices_np, os.path.dirname(cfg.TEST.WEIGHT), dataset.test_tracks)
def main(): parser = argparse.ArgumentParser(description="ReID Baseline Inference") parser.add_argument("--config_file", default="./configs/debug.yml", help="path to config file", type=str) parser.add_argument("opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER) args = parser.parse_args() num_gpus = int( os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1 if args.config_file != "": cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() output_dir = cfg.OUTPUT_DIR if output_dir and not os.path.exists(output_dir): mkdir(output_dir) logger = setup_logger("reid_baseline", output_dir, 0) logger.info("Using {} GPUS".format(num_gpus)) logger.info(args) if args.config_file != "": logger.info("Loaded configuration file {}".format(args.config_file)) # with open(args.config_file, 'r') as cf: # config_str = "\n" + cf.read() # logger.info(config_str) logger.info("Running with config:\n{}".format(cfg)) if cfg.MODEL.DEVICE == "cuda": os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID cudnn.benchmark = True train_loader, val_loader, num_query, num_classes, dataset = make_data_loader( cfg) model = build_model(cfg, num_classes) model.load_param(cfg.TEST.WEIGHT) imgs, feats = inference_to_get_feats(cfg, model, val_loader, num_query, dataset) import numpy as np np.save('feats_extract.npy', feats.cpu().detach().numpy()) print(imgs[:50])
def main(): parser = argparse.ArgumentParser(description="ReID Baseline Inference") parser.add_argument("--config_file", default="./configs/debug.yml", help="path to config file", type=str) parser.add_argument("opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER) args = parser.parse_args() num_gpus = int( os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1 if args.config_file != "": cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() output_dir = cfg.OUTPUT_DIR if output_dir and not os.path.exists(output_dir): mkdir(output_dir) logger = setup_logger("reid_baseline", output_dir, 0) logger.info("Using {} GPUS".format(num_gpus)) logger.info(args) if args.config_file != "": logger.info("Loaded configuration file {}".format(args.config_file)) logger.info("Running with config:\n{}".format(cfg)) if cfg.MODEL.DEVICE == "cuda": os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID cudnn.benchmark = True train_loader, val_loader, num_query, num_classes, dataset = make_data_loader( cfg) model = build_model(cfg, num_classes) model.load_param(cfg.TEST.WEIGHT) results = [] out_dir = os.path.join('AIC_crop') results += vis_actmap(model, cfg, train_loader, out_dir) results += vis_actmap(model, cfg, val_loader, out_dir) with open(os.path.join(out_dir, 'detection.json'), 'w') as f: json.dump(results, f)
def main(): parser = argparse.ArgumentParser() parser.add_argument('--config-file', default='', help='Path to config file', type=str) parser.add_argument('opts', help="Modify config options using the command line", default=None, nargs=argparse.REMAINDER) args = parser.parse_args() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() train_net(cfg)
def main(): parser = argparse.ArgumentParser(description="ReID Baseline Training") parser.add_argument( # "--config_file", default="/media/bi/Data/Mycode/car_demo/AICity2020-VOC-ReID-7c453723e6e9179d175772921f93441cfa621dc1/configs/aicity20.yml", help="path to config file", type=str "--config_file", default="data/veri.yml", help="path to config file", type=str) parser.add_argument("opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER) args = parser.parse_args() num_gpus = int( os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1 if args.config_file != "": cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() output_dir = cfg.OUTPUT_DIR if output_dir and not os.path.exists(output_dir): os.makedirs(output_dir) logger = setup_logger("reid_baseline", output_dir, 0) logger.info("Using {} GPUS".format(num_gpus)) logger.info(args) if args.config_file != "": logger.info("Loaded configuration file {}".format(args.config_file)) # with open(args.config_file, 'r') as cf: # config_str = "\n" + cf.read() # logger.info(config_str) logger.info("Running with config:\n{}".format(cfg)) if cfg.MODEL.DEVICE == "cuda": os.environ[ 'CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID # new add by gu cudnn.benchmark = True train(cfg)
def main(): ''' parse config file ''' parser = argparse.ArgumentParser(description="Graph Reasoning Machine for Visual Question Answering") parser.add_argument("--config-file", default="configs/baseline_res101.yaml") parser.add_argument("--local_rank", type=int, default=0) parser.add_argument("--resume", type=int, default=0) parser.add_argument("--inference", action='store_true') parser.add_argument("--instance", type=int, default=-1) parser.add_argument("--use_freq_prior", action='store_true') parser.add_argument("--visualize", action='store_true') parser.add_argument("--algorithm", type=str, default='sg_baseline') args = parser.parse_args() num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1 args.distributed = num_gpus > 1 if args.distributed: torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group( backend="nccl", init_method="env://" ) synchronize() cfg.merge_from_file(args.config_file) cfg.resume = args.resume cfg.instance = args.instance cfg.inference = args.inference cfg.MODEL.USE_FREQ_PRIOR = args.use_freq_prior cfg.MODEL.ALGORITHM = args.algorithm cfg.freeze() if not os.path.exists("logs") and get_rank() == 0: os.mkdir("logs") logger = setup_logger("scene_graph_generation", "logs", get_rank()) logger.info(args) logger.info("Loaded configuration file {}".format(args.config_file)) output_config_path = os.path.join("logs", 'config.yml') logger.info("Saving config into: {}".format(output_config_path)) save_config(cfg, output_config_path) if not args.inference: model = train(cfg, args) else: test(cfg, args)
def main(): parser = argparse.ArgumentParser() parser.add_argument('--config-file', default='', help='Path to config file', type=str) parser.add_argument('--test', help='Test or train model', action='store_true') parser.add_argument('opts', help="Modify config options using the command line", default=None, nargs=argparse.REMAINDER) args = parser.parse_args() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() if args.test: test_net(cfg) else: train_net(cfg)
def main(): parser = argparse.ArgumentParser( description="PyTorch Object Detection Training") parser.add_argument( "--config-file", default="data/occlusion_net_train.yaml", metavar="FILE", help="path to config file", type=str, ) parser.add_argument("--local_rank", type=int, default=0) parser.add_argument( "--cometml-tag", dest="cometml_tag", default="occlusion-net", ) parser.add_argument( "--skip-test", dest="skip_test", help="Do not test the final model", action="store_true", ) parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args() num_gpus = int( os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1 args.distributed = num_gpus > 1 if args.distributed: torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group(backend="nccl", init_method="env://") synchronize() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() output_dir = cfg.OUTPUT_DIR if output_dir: mkdir(output_dir) logger = setup_logger("maskrcnn_benchmark", output_dir, get_rank()) logger.info("Using {} GPUs".format(num_gpus)) logger.info(args) logger.info("Collecting env info (might take some time)") logger.info("\n" + collect_env_info()) logger.info("Loaded configuration file {}".format(args.config_file)) with open(args.config_file, "r") as cf: config_str = "\n" + cf.read() logger.info(config_str) logger.info("Running with config:\n{}".format(cfg)) #model = train(cfg, args.local_rank, args.distributed) model = build_detection_model(cfg) device = torch.device(cfg.MODEL.DEVICE) model.to(device).eval() if not args.skip_test: run_test(cfg, model, args.distributed)
def main(): print('Load config...') args = parse_args() update_config(cfg, args) # cfg.defrost() # cfg.RANK = args.ranka # cfg.freeze() # device allocation print('Set device...') #print(cfg.GPUS) #os.environ["CUDA_VISIBLE_DEVICES"] = cfg.GPUS #device = torch.device('cuda') torch.cuda.set_device(cfg.GPUS[0]) device = torch.device('cuda:' + str(cfg.GPUS[0])) print("Build dataloader ...") # load texture if cfg.DATASET.TEX_PATH: texture_init = cv2.cvtColor(cv2.imread(cfg.DATASET.TEX_PATH), cv2.COLOR_BGR2RGB) texture_init_resize = cv2.resize( texture_init, (cfg.MODEL.TEX_MAPPER.NUM_SIZE, cfg.MODEL.TEX_MAPPER.NUM_SIZE), interpolation=cv2.INTER_AREA).astype(np.float32) / 255.0 texture_init_use = torch.from_numpy(texture_init_resize).to(device) # dataset for training views view_dataset = dataio.ViewDataset( cfg=cfg, root_dir=cfg.DATASET.ROOT, calib_path=cfg.DATASET.CALIB_PATH, calib_format=cfg.DATASET.CALIB_FORMAT, sampling_pattern=cfg.TRAIN.SAMPLING_PATTERN, precomp_high_dir=cfg.DATASET.PRECOMP_DIR, precomp_low_dir=cfg.DATASET.PRECOMP_DIR, preset_uv_path=cfg.DATASET.UV_PATH, ) # dataset for validation views view_val_dataset = dataio.ViewDataset( cfg=cfg, root_dir=cfg.DATASET.ROOT, calib_path=cfg.DATASET.CALIB_PATH, calib_format=cfg.DATASET.CALIB_FORMAT, sampling_pattern=cfg.TRAIN.SAMPLING_PATTERN_VAL, precomp_high_dir=cfg.DATASET.PRECOMP_DIR, precomp_low_dir=cfg.DATASET.PRECOMP_DIR, ) num_view_val = len(view_val_dataset) print('Build Network...') # Rasterizer cur_obj_path = '' if not cfg.DATASET.LOAD_PRECOMPUTE: view_data = view_dataset.read_view(0) cur_obj_path = view_data['obj_path'] frame_idx = view_data['f_idx'] obj_data = view_dataset.objs[frame_idx] rasterizer = network.Rasterizer( cfg, obj_fp=cur_obj_path, img_size=cfg.DATASET.OUTPUT_SIZE[0], camera_mode=cfg.DATASET.CAM_MODE, obj_data=obj_data, # preset_uv_path = cfg.DATASET.UV_PATH, global_RT=view_dataset.global_RT) # texture mapper texture_mapper = network.TextureMapper( texture_size=cfg.MODEL.TEX_MAPPER.NUM_SIZE, texture_num_ch=cfg.MODEL.TEX_MAPPER.NUM_CHANNELS, mipmap_level=cfg.MODEL.TEX_MAPPER.MIPMAP_LEVEL, apply_sh=cfg.MODEL.TEX_MAPPER.SH_BASIS) # render net render_net = network.RenderingNet( nf0=cfg.MODEL.RENDER_NET.NF0, in_channels=cfg.MODEL.TEX_MAPPER.NUM_CHANNELS, out_channels=3, num_down_unet=5, use_gcn=False) # interpolater interpolater = network.Interpolater() # L1 loss criterionL1 = nn.L1Loss(reduction='mean').to(device) # Optimizer optimizerG = torch.optim.Adam(list(texture_mapper.parameters()) + list(render_net.parameters()), lr=cfg.TRAIN.LR) print('Loading Model...') iter = 0 dir_name = os.path.join(datetime.datetime.now().strftime('%m-%d') + '_' + datetime.datetime.now().strftime('%H-%M-%S') + '_' + cfg.TRAIN.SAMPLING_PATTERN + '_' + cfg.DATASET.ROOT.strip('/').split('/')[-1]) if cfg.TRAIN.EXP_NAME is not '': dir_name += '_' + cfg.TRAIN.EXP_NAME if cfg.AUTO_RESUME: checkpoint_path = '' if cfg.TRAIN.RESUME and cfg.TRAIN.CHECKPOINT: checkpoint_path = cfg.TRAIN.CHECKPOINT dir_name = cfg.TRAIN.CHECKPOINT_DIR nums = [ int(s) for s in cfg.TRAIN.CHECKPOINT_NAME.split('_') if s.isdigit() ] cfg.defrost() cfg.TRAIN.BEGIN_EPOCH = nums[0] + 1 cfg.freeze() iter = nums[1] + 1 elif cfg.MODEL.PRETRAINED: checkpoint_path = cfg.MODEL.PRETRAIN if checkpoint_path: print(' Checkpoint_path : %s' % (checkpoint_path)) util.custom_load([texture_mapper, render_net], ['texture_mapper', 'render_net'], checkpoint_path) else: print(' Not load params. ') texture_mapper.to(device) render_net.to(device) interpolater.to(device) rasterizer.to(device) texture_mapper_module = texture_mapper render_net_module = render_net # use multi-GPU if len(cfg.GPUS) > 1: texture_mapper = nn.DataParallel(texture_mapper, device_ids=cfg.GPUS) render_net = nn.DataParallel(render_net, device_ids=cfg.GPUS) interpolater = nn.DataParallel(interpolater, device_ids=cfg.GPUS) rasterizer = nn.DataParallel(rasterizer, device_ids=cfg.GPUS) rasterizer = rasterizer.module # set to training mode texture_mapper.train() render_net.train() interpolater.train() rasterizer.eval() # not train now part_list = [texture_mapper_module, render_net_module] # collect all networks part_name_list = ['texture_mapper', 'render_net'] print("*" * 100) print("Number of generator parameters:") cfg.defrost() cfg.MODEL.TEX_MAPPER.NUM_PARAMS = util.print_network(texture_mapper).item() cfg.MODEL.RENDER_NET.NUM_PARAMS = util.print_network(render_net).item() cfg.freeze() print("*" * 100) print("Setup Log ...") log_dir = os.path.join(cfg.LOG.LOGGING_ROOT, dir_name) data_util.cond_mkdir(log_dir) val_out_dir = os.path.join(log_dir, 'val_out') val_gt_dir = os.path.join(log_dir, 'val_gt') val_err_dir = os.path.join(log_dir, 'val_err') data_util.cond_mkdir(val_out_dir) data_util.cond_mkdir(val_gt_dir) data_util.cond_mkdir(val_err_dir) util.custom_copy(args.cfg, os.path.join(log_dir, cfg.LOG.CFG_NAME)) print('Start buffering data for training and validation...') view_dataloader = DataLoader(view_dataset, batch_size=cfg.TRAIN.BATCH_SIZE, shuffle=cfg.TRAIN.SHUFFLE, num_workers=8) view_val_dataloader = DataLoader(view_val_dataset, batch_size=cfg.TRAIN.BATCH_SIZE, shuffle=False, num_workers=8) #view_dataset.buffer_all() #view_val_dataset.buffer_all() # Save all command line arguments into a txt file in the logging directory for later referene. writer = SummaryWriter(log_dir) # iter = cfg.TRAIN.BEGIN_EPOCH * len(view_dataset) # pre model is batch-1 print('Begin training...') # init value val_log_batch_id = 0 first_val = True img_h, img_w = cfg.DATASET.OUTPUT_SIZE for epoch in range(cfg.TRAIN.BEGIN_EPOCH, cfg.TRAIN.END_EPOCH): for view_trgt in view_dataloader: start = time.time() # get image img_gt = [] img_gt.append(view_trgt[0]['img_gt'].to(device)) ROI = view_trgt[0]['ROI'].to(device) # get uvmap alpha uv_map = [] alpha_map = [] if not cfg.DATASET.LOAD_PRECOMPUTE: # raster module frame_idxs = view_trgt[0]['f_idx'].numpy() for batch_idx, frame_idx in enumerate(frame_idxs): obj_path = view_trgt[0]['obj_path'][batch_idx] if cur_obj_path != obj_path: cur_obj_path = obj_path obj_data = view_dataset.objs[frame_idx] rasterizer.update_vs(obj_data['v_attr']) proj = view_trgt[0]['proj'].to(device)[batch_idx, ...] pose = view_trgt[0]['pose'].to(device)[batch_idx, ...] dist_coeffs = view_trgt[0]['dist_coeffs'].to(device)[ batch_idx, ...] uv_map_single, alpha_map_single, _, _, _, _, _, _, _, _, _, _, _, _ = \ rasterizer(proj = proj[None, ...], pose = pose[None, ...], dist_coeffs = dist_coeffs[None, ...], offset = None, scale = None, ) uv_map.append(uv_map_single[0, ...].clone().detach()) alpha_map.append(alpha_map_single[0, ...].clone().detach()) # fix alpha map uv_map = torch.stack(uv_map, dim=0) alpha_map = torch.stack(alpha_map, dim=0)[:, None, :, :] # alpha_map = alpha_map * torch.tensor(img_gt[0][:,0,:,:][:,None,:,:] <= (2.0 * 255)).permute(0,2,1,3).to(alpha_map.dtype).to(alpha_map.device) # check per iter image for batch_idx, frame_idx in enumerate(frame_idxs): if cfg.DEBUG.SAVE_TRANSFORMED_IMG: save_dir_img_gt = './Debug/image_mask' save_path_img_gt = os.path.join( save_dir_img_gt, '%06d_%03d.png' % (iter, frame_idx)) cv2.imwrite( save_path_img_gt, cv2.cvtColor( img_gt[0][batch_idx, ...].cpu().detach().numpy( ).transpose(1, 2, 0) * 255.0, cv2.COLOR_RGB2BGR)) #cv2.imwrite(os.path.join(save_dir_img_gt, '%03d_'%frame_idx + img_fn), cv2.cvtColor(img_gt*255.0, cv2.COLOR_BGR2RGB)) print(' Save img: ' + save_path_img_gt) if cfg.DEBUG.SAVE_TRANSFORMED_MASK: save_alpha_map = alpha_map.permute( 0, 2, 3, 1).cpu().detach().numpy() save_dir_mask = './Debug/image_mask' save_path_mask = os.path.join( save_dir_mask, '%06d_%03d_mask.png' % (iter, frame_idx)) cv2.imwrite(save_path_mask, save_alpha_map[batch_idx, ...] * 255.0) print(' Save mask: ' + save_path_mask) else: # get view data uv_map = view_trgt[0]['uv_map'].to(device) # [N, H, W, 2] # sh_basis_map = view_trgt[0]['sh_basis_map'].to(device) # [N, H, W, 9] alpha_map = view_trgt[0]['alpha_map'][:, None, :, :].to( device) # [N, 1, H, W] # sample texture # neural_img = texture_mapper(uv_map, sh_basis_map) neural_img = texture_mapper(uv_map) # rendering net outputs = render_net(neural_img, None) img_max_val = 2.0 outputs = (outputs * 0.5 + 0.5) * img_max_val # map to [0, img_max_val] if type(outputs) is not list: outputs = [outputs] # # We don't enforce a loss on the outermost 5 pixels to alleviate boundary errors, also weight loss by alpha # alpha_map_central = alpha_map[:, :, 5:-5, 5:-5] # for i in range(len(view_trgt)): # outputs[i] = outputs[i][:, :, 5:-5, 5:-5] * alpha_map_central # img_gt[i] = img_gt[i][:, :, 5:-5, 5:-5] * alpha_map_central # ignore loss outside ROI for i in range(len(view_trgt)): outputs[i] = outputs[i] * ROI * alpha_map img_gt[i] = img_gt[i] * ROI * alpha_map # loss on final image loss_rn = list() for idx in range(len(view_trgt)): loss_rn.append( criterionL1(outputs[idx].contiguous().view(-1).float(), img_gt[idx].contiguous().view(-1).float())) loss_rn = torch.stack(loss_rn, dim=0).mean() # total loss for generator loss_g = loss_rn optimizerG.zero_grad() loss_g.backward() optimizerG.step() # error metrics with torch.no_grad(): err_metrics_batch_i = metric.compute_err_metrics_batch( outputs[0] * 255.0, img_gt[0] * 255.0, alpha_map, compute_ssim=False) # err_metrics_batch_i = metric.compute_err_metrics_batch(outputs[0] * 255.0, img_gt[0] * 255.0, alpha_map_central, compute_ssim = False) # tensorboard scalar logs of training data writer.add_scalar("loss_g", loss_g, iter) writer.add_scalar("loss_rn", loss_rn, iter) writer.add_scalar("final_mae_valid", err_metrics_batch_i['mae_valid_mean'], iter) writer.add_scalar("final_psnr_valid", err_metrics_batch_i['psnr_valid_mean'], iter) end = time.time() print( "Iter %07d Epoch %03d loss_g %0.4f mae_valid %0.4f psnr_valid %0.4f t_total %0.4f" % (iter, epoch, loss_g, err_metrics_batch_i['mae_valid_mean'], err_metrics_batch_i['psnr_valid_mean'], end - start)) # tensorboard figure logs of training data if not iter % cfg.LOG.PRINT_FREQ: output_final_vs_gt = [] for i in range(len(view_trgt)): output_final_vs_gt.append(outputs[i].clamp(min=0., max=1.)) output_final_vs_gt.append(img_gt[i].clamp(min=0., max=1.)) output_final_vs_gt.append( (outputs[i] - img_gt[i]).abs().clamp(min=0., max=1.)) output_final_vs_gt = torch.cat(output_final_vs_gt, dim=0) raster_uv_maps = torch.cat( ( uv_map.permute(0, 3, 1, 2), # N H W 2 -> N 2 H W torch.zeros(uv_map.shape[0], 1, img_h, img_w, dtype=uv_map.dtype, device=uv_map.device)), dim=1) writer.add_image( "raster_uv_vis", torchvision.utils.make_grid( raster_uv_maps, nrow=raster_uv_maps[0].shape[0], range=(0, 1), scale_each=False, normalize=False).cpu().detach().numpy() [::-1, :, :], # uv0 -> 0vu (rgb) iter) writer.add_image( "output_final_vs_gt", torchvision.utils.make_grid( output_final_vs_gt, nrow=outputs[0].shape[0], # 3 range=(0, 1), scale_each=False, normalize=False).cpu().detach().numpy(), iter) # validation if not iter % cfg.TRAIN.VAL_FREQ: start_val = time.time() with torch.no_grad(): # error metrics err_metrics_val = {} err_metrics_val['mae_valid'] = [] err_metrics_val['mse_valid'] = [] err_metrics_val['psnr_valid'] = [] err_metrics_val['ssim_valid'] = [] # loop over batches batch_id = 0 for view_val_trgt in view_val_dataloader: start_val_i = time.time() # get image img_gt = [] img_gt.append(view_val_trgt[0]['img_gt'].to(device)) ROI = view_val_trgt[0]['ROI'].to(device) # get uvmap alpha uv_map = [] alpha_map = [] if not cfg.DATASET.LOAD_PRECOMPUTE: # build raster module frame_idxs = view_val_trgt[0]['f_idx'].numpy() for batch_idx, frame_idx in enumerate(frame_idxs): obj_path = view_val_trgt[0]['obj_path'][ batch_idx] if cur_obj_path != obj_path: cur_obj_path = obj_path obj_data = view_val_dataset.objs[frame_idx] rasterizer.update_vs(obj_data['v_attr']) proj = view_val_trgt[0]['proj'].to(device)[ batch_idx, ...] pose = view_val_trgt[0]['pose'].to(device)[ batch_idx, ...] dist_coeffs = view_val_trgt[0][ 'dist_coeffs'].to(device)[batch_idx, ...] uv_map_single, alpha_map_single, _, _, _, _, _, _, _, _, _, _, _, _ = \ rasterizer(proj = proj[None, ...], pose = pose[None, ...], dist_coeffs = dist_coeffs[None, ...], offset = None, scale = None, ) uv_map.append( uv_map_single[0, ...].clone().detach()) alpha_map.append( alpha_map_single[0, ...].clone().detach()) # fix alpha map uv_map = torch.stack(uv_map, dim=0) alpha_map = torch.stack(alpha_map, dim=0)[:, None, :, :] # alpha_map = alpha_map * torch.tensor(img_gt[0][:,0,:,:][:,None,:,:] <= (2.0 * 255)).permute(0,2,1,3).to(alpha_map.dtype).to(alpha_map.device) else: uv_map = view_val_trgt[0]['uv_map'].to( device) # [N, H, W, 2] # sh_basis_map = view_val_trgt[0]['sh_basis_map'].to(device) # [N, H, W, 9] alpha_map = view_val_trgt[0][ 'alpha_map'][:, None, :, :].to( device) # [N, 1, H, W] view_idx = view_val_trgt[0]['idx'] num_view = len(view_val_trgt) img_gt = [] for i in range(num_view): img_gt.append( view_val_trgt[i]['img_gt'].to(device)) # sample texture # neural_img = texture_mapper(uv_map, sh_basis_map) neural_img = texture_mapper(uv_map) # rendering net outputs = render_net(neural_img, None) img_max_val = 2.0 outputs = (outputs * 0.5 + 0.5 ) * img_max_val # map to [0, img_max_val] if type(outputs) is not list: outputs = [outputs] # apply alpha and ROI for i in range(num_view): outputs[i] = outputs[i] * alpha_map * ROI img_gt[i] = img_gt[i] * alpha_map * ROI # tensorboard figure logs of validation data if batch_id == val_log_batch_id: output_final_vs_gt = [] for i in range(num_view): output_final_vs_gt.append(outputs[i].clamp( min=0., max=1.)) output_final_vs_gt.append(img_gt[i].clamp( min=0., max=1.)) output_final_vs_gt.append( (outputs[i] - img_gt[i]).abs().clamp( min=0., max=1.)) output_final_vs_gt = torch.cat(output_final_vs_gt, dim=0) writer.add_image( "output_final_vs_gt_val", torchvision.utils.make_grid( output_final_vs_gt, nrow=outputs[0].shape[0], # 3 range=(0, 1), scale_each=False, normalize=False).cpu().detach().numpy(), iter) # error metrics err_metrics_batch_i_final = metric.compute_err_metrics_batch( outputs[0] * 255.0, img_gt[0] * 255.0, alpha_map, compute_ssim=True) batch_size = view_idx.shape[0] for i in range(batch_size): for key in list(err_metrics_val.keys()): if key in err_metrics_batch_i_final.keys(): err_metrics_val[key].append( err_metrics_batch_i_final[key][i]) # save images for i in range(batch_size): cv2.imwrite( os.path.join( val_out_dir, str(iter).zfill(8) + '_' + str(view_idx[i].cpu().detach().numpy( )).zfill(5) + '.png'), outputs[0][i, :].permute( (1, 2, 0)).cpu().detach().numpy()[:, :, ::-1] * 255.) cv2.imwrite( os.path.join( val_err_dir, str(iter).zfill(8) + '_' + str(view_idx[i].cpu().detach().numpy( )).zfill(5) + '.png'), (outputs[0] - img_gt[0]).abs().clamp( min=0., max=1.)[i, :].permute( (1, 2, 0)).cpu().detach().numpy()[:, :, ::-1] * 255.) if first_val: cv2.imwrite( os.path.join( val_gt_dir, str(view_idx[i].cpu().detach().numpy() ).zfill(5) + '.png'), img_gt[0][i, :].permute( (1, 2, 0)).cpu().detach().numpy()[:, :, ::-1] * 255.) end_val_i = time.time() print( "Val batch %03d mae_valid %0.4f psnr_valid %0.4f ssim_valid %0.4f t_total %0.4f" % (batch_id, err_metrics_batch_i_final['mae_valid_mean'], err_metrics_batch_i_final['psnr_valid_mean'], err_metrics_batch_i_final['ssim_valid_mean'], end_val_i - start_val_i)) batch_id += 1 for key in list(err_metrics_val.keys()): if err_metrics_val[key]: err_metrics_val[key] = np.vstack( err_metrics_val[key]) err_metrics_val[ key + '_mean'] = err_metrics_val[key].mean() else: err_metrics_val[key + '_mean'] = np.nan # tensorboard scalar logs of validation data writer.add_scalar("final_mae_valid_val", err_metrics_val['mae_valid_mean'], iter) writer.add_scalar("final_psnr_valid_val", err_metrics_val['psnr_valid_mean'], iter) writer.add_scalar("final_ssim_valid_val", err_metrics_val['ssim_valid_mean'], iter) first_val = False val_log_batch_id = (val_log_batch_id + 1) % batch_id end_val = time.time() print( "Val mae_valid %0.4f psnr_valid %0.4f ssim_valid %0.4f t_total %0.4f" % (err_metrics_val['mae_valid_mean'], err_metrics_val['psnr_valid_mean'], err_metrics_val['ssim_valid_mean'], end_val - start_val)) iter += 1 if iter % cfg.LOG.CHECKPOINT_FREQ == 0: util.custom_save( os.path.join(log_dir, 'model_epoch_%d_iter_%s_.pth' % (epoch, iter)), part_list, part_name_list) util.custom_save( os.path.join(log_dir, 'model_epoch_%d_iter_%s_.pth' % (epoch, iter)), part_list, part_name_list)
"--config-file", default="", metavar="FILE", help="Path to config file", type=str, ) parser.add_argument( "--skip-test", dest="skip_test", help="Do not test the final model", action="store_true", ) parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER ) parser.add_argument( "--type", type=str, ) args = parser.parse_args() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() globals()["run_" + args.type](cfg)
def update_config(cfg, yamlfilename): cfg.defrost() cfg.merge_from_file(yamlfilename) cfg.TEST.MODEL_FILE = osp.join(cfg.DATA_DIR, cfg.TEST.MODEL_FILE) cfg.freeze()
def main(): parser = argparse.ArgumentParser(description="ReID Baseline Inference") parser.add_argument( "--config_file", default="./configs/submit.yml", help="path to config file", type=str ) parser.add_argument("opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER) parser.add_argument( # "--config_file", default="/media/bi/Data/Mycode/car_demo/AICity2020-VOC-ReID-7c453723e6e9179d175772921f93441cfa621dc1/configs/aicity20.yml", help="path to config file", type=str "--pretrain_model_path", default="./dianshang/crop_half_model.pth", help="path to config file", type=str ) parser.add_argument( # "--config_file", default="/media/bi/Data/Mycode/car_demo/AICity2020-VOC-ReID-7c453723e6e9179d175772921f93441cfa621dc1/configs/aicity20.yml", help="path to config file", type=str "--crop_path", default=" ", help="path to config file", type=str ) args = parser.parse_args() num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1 if args.config_file != "": cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.TEST.WEIGHT = args.pretrain_model_path cfg.freeze() output_dir = cfg.OUTPUT_DIR if output_dir and not os.path.exists(output_dir): mkdir(output_dir) logger = setup_logger("reid_baseline", output_dir, 0) logger.info("Using {} GPUS".format(num_gpus)) logger.info(args) if args.config_file != "": logger.info("Loaded configuration file {}".format(args.config_file)) logger.info("Running with config:\n{}".format(cfg)) if cfg.MODEL.DEVICE == "cuda": os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID cudnn.benchmark = True train_loader, val_loader, num_query, num_classes, dataset = make_data_loader(cfg,shuffle_train= False) model = build_model(cfg, num_classes) print("load pretained weight",cfg.TEST.WEIGHT) # try: print("加载单卡训练权重") model.load_param(cfg.TEST.WEIGHT) results = [] # out_dir = '/home/lab3/bi/0716/Veri/ai_city/tools/output/crop/crop_query_new' # out_dir = '/home/lab3/bi/0716/Veri/ai_city/tools/output/crop/crop_gallery_new' # out_dir = '/home/lab3/bi/0716/Veri/ai_city/tools/output/crop/crop_query' out_dir = args.crop_path if os.path.exists(os.path.join(out_dir,'crop_train')): print("文件夹已存在") else: os.makedirs(os.path.join(out_dir,'crop_train')) os.makedirs(os.path.join(out_dir,'crop_query')) os.makedirs(os.path.join(out_dir,'crop_gallery')) crop_train,crop_query,crop_gallery= vis_actmap(model, cfg, train_loader, out_dir) pickle.dump(crop_train, open(os.path.join(out_dir, 'train_crop_img_add.pkl'), 'wb')) crop_train,crop_query,crop_gallery= vis_actmap(model, cfg, val_loader, out_dir) pickle.dump((crop_query, crop_gallery), open(os.path.join(out_dir, 'test_crop_img_add.pkl'), 'wb'))