def train_joint(config, output_dir, args): assert 'train_iter' in config # config # from utils.utils import pltImshow # from utils.utils import saveImg print("Start") torch.set_default_tensor_type(torch.FloatTensor) task = config['data']['dataset'] print("Start 1") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") logging.info('train on device: %s', device) print("Start 2:", config['front_end_model']) with open(os.path.join(output_dir, 'config.yml'), 'w') as f: yaml.dump(config, f, default_flow_style=False) # writer = SummaryWriter(getWriterPath(task=args.command, date=True)) writer = SummaryWriter( getWriterPath(task=args.command, exper_name=args.exper_name, date=True)) print("Start 3") ## save data save_path = get_save_path(output_dir) print("Save path: ", save_path) # data loading # data = dataLoader(config, dataset='syn', warp_input=True) print("Start 4: ", task) data = dataLoader(config, dataset=task, warp_input=True) print("Start 5") train_loader, val_loader = data['train_loader'], data['val_loader'] print("Mid ") datasize(train_loader, config, tag='train') datasize(val_loader, config, tag='val') # init the training agent using config file # from train_model_frontend import Train_model_frontend from utils.loader import get_module train_model_frontend = get_module('', config['front_end_model']) print("Front end is: ", train_model_frontend) train_agent = train_model_frontend(config, save_path=save_path, device=device) # writer from tensorboard train_agent.writer = writer # feed the data into the agent train_agent.train_loader = train_loader train_agent.val_loader = val_loader # load model initiates the model and load the pretrained model (if any) train_agent.loadModel() train_agent.dataParallel() try: # train function takes care of training and evaluation train_agent.train() except KeyboardInterrupt: print("press ctrl + c, save model!") train_agent.saveModel() pass
def export_descriptor(config, output_dir, args): """ # input 2 images, output keypoints and correspondence save prediction: pred: 'image': np(320,240) 'prob' (keypoints): np (N1, 2) 'desc': np (N2, 256) 'warped_image': np(320,240) 'warped_prob' (keypoints): np (N2, 2) 'warped_desc': np (N2, 256) 'homography': np (3,3) 'matches': np [N3, 4] """ from utils.loader import get_save_path from utils.var_dim import squeezeToNumpy # basic settings device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") logging.info("train on device: %s", device) with open(os.path.join(output_dir, "config.yml"), "w") as f: yaml.dump(config, f, default_flow_style=False) writer = SummaryWriter(getWriterPath(task=args.command, date=True)) save_path = get_save_path(output_dir) save_output = save_path / "../predictions" os.makedirs(save_output, exist_ok=True) ## parameters outputMatches = True subpixel = config["model"]["subpixel"]["enable"] patch_size = config["model"]["subpixel"]["patch_size"] # data loading from utils.loader import dataLoader_test as dataLoader task = config["data"]["dataset"] data = dataLoader(config, dataset=task) test_set, test_loader = data["test_set"], data["test_loader"] from utils.print_tool import datasize datasize(test_loader, config, tag="test") # model loading from utils.loader import get_module Val_model_heatmap = get_module("", config["front_end_model"]) ## load pretrained val_agent = Val_model_heatmap(config["model"], device=device) val_agent.loadModel() ## tracker tracker = PointTracker(max_length=2, nn_thresh=val_agent.nn_thresh) ###### check!!! count = 0 for i, sample in tqdm(enumerate(test_loader)): img_0, img_1 = sample["image"], sample["warped_image"] # first image, no matches # img = img_0 def get_pts_desc_from_agent(val_agent, img, device="cpu"): """ pts: list [numpy (3, N)] desc: list [numpy (256, N)] """ heatmap_batch = val_agent.run( img.to(device) ) # heatmap: numpy [batch, 1, H, W] # heatmap to pts pts = val_agent.heatmap_to_pts() # print("pts: ", pts) if subpixel: pts = val_agent.soft_argmax_points(pts, patch_size=patch_size) # heatmap, pts to desc desc_sparse = val_agent.desc_to_sparseDesc() # print("pts[0]: ", pts[0].shape, ", desc_sparse[0]: ", desc_sparse[0].shape) # print("pts[0]: ", pts[0].shape) outs = {"pts": pts[0], "desc": desc_sparse[0]} return outs def transpose_np_dict(outs): for entry in list(outs): outs[entry] = outs[entry].transpose() outs = get_pts_desc_from_agent(val_agent, img_0, device=device) pts, desc = outs["pts"], outs["desc"] # pts: np [3, N] if outputMatches == True: tracker.update(pts, desc) # save keypoints pred = {"image": squeezeToNumpy(img_0)} pred.update({"prob": pts.transpose(), "desc": desc.transpose()}) # second image, output matches outs = get_pts_desc_from_agent(val_agent, img_1, device=device) pts, desc = outs["pts"], outs["desc"] if outputMatches == True: tracker.update(pts, desc) pred.update({"warped_image": squeezeToNumpy(img_1)}) # print("total points: ", pts.shape) pred.update( { "warped_prob": pts.transpose(), "warped_desc": desc.transpose(), "homography": squeezeToNumpy(sample["homography"]), } ) if outputMatches == True: matches = tracker.get_matches() print("matches: ", matches.transpose().shape) pred.update({"matches": matches.transpose()}) print("pts: ", pts.shape, ", desc: ", desc.shape) # clean last descriptor tracker.clear_desc() filename = str(count) path = Path(save_output, "{}.npz".format(filename)) np.savez_compressed(path, **pred) # print("save: ", path) count += 1 print("output pairs: ", count)
def export_detector_homoAdapt_gpu(config, output_dir, args): """ input 1 images, output pseudo ground truth by homography adaptation. Save labels: pred: 'prob' (keypoints): np (N1, 3) """ from utils.utils import pltImshow from utils.utils import saveImg from utils.draw import draw_keypoints # basic setting task = config["data"]["dataset"] export_task = config["data"]["export_folder"] device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") logging.info("train on device: %s", device) with open(os.path.join(output_dir, "config.yml"), "w") as f: yaml.dump(config, f, default_flow_style=False) writer = SummaryWriter( getWriterPath(task=args.command, exper_name=args.exper_name, date=True) ) ## parameters nms_dist = config["model"]["nms"] # 4 top_k = config["model"]["top_k"] homoAdapt_iter = config["data"]["homography_adaptation"]["num"] conf_thresh = config["model"]["detection_threshold"] nn_thresh = 0.7 outputMatches = True count = 0 max_length = 5 output_images = args.outputImg check_exist = True ## save data save_path = Path(output_dir) save_output = save_path save_output = save_output / "predictions" / export_task save_path = save_path / "checkpoints" logging.info("=> will save everything to {}".format(save_path)) os.makedirs(save_path, exist_ok=True) os.makedirs(save_output, exist_ok=True) # data loading from utils.loader import dataLoader_test as dataLoader data = dataLoader(config, dataset=task, export_task=export_task) print("Data is: ",data) test_set, test_loader = data["test_set"], data["test_loader"] print("Size test: ",len(test_set)) print("Size loader: ",len(test_loader)) # model loading ## load pretrained try: path = config["pretrained"] print("==> Loading pre-trained network.") print("path: ", path) # This class runs the SuperPoint network and processes its outputs. fe = SuperPointFrontend_torch( config=config, weights_path=path, nms_dist=nms_dist, conf_thresh=conf_thresh, nn_thresh=nn_thresh, cuda=False, device=device, ) print("==> Successfully loaded pre-trained network.") fe.net_parallel() print(path) # save to files save_file = save_output / "export.txt" with open(save_file, "a") as myfile: myfile.write("load model: " + path + "\n") except Exception: print(f"load model: {path} failed! ") raise def load_as_float(path): return imread(path).astype(np.float32) / 255 print("Tracker") tracker = PointTracker(max_length, nn_thresh=fe.nn_thresh) with open(save_file, "a") as myfile: myfile.write("homography adaptation: " + str(homoAdapt_iter) + "\n") print("Load save file") ''' print(len(test_loader)) for i,sample in enumerate(test_loader): print("Hello world") print("Img: ",sample["image"].size()) print("Name: ",test_set[i]["name"]) print("valid mask: ",test_set[i]["valid_mask"].size()) print("valid img_2D: ",test_set[i]["image_2D"].size()) print("valid mask: ",test_set[i]["valid_mask"].size()) print("homograpgy: ",test_set[i]["homographies"].size()) print("inverse: ",test_set[i]["inv_homographies"].size()) print("scene name: ",test_set[i]["scene_name"]) print() ''' ## loop through all images for i, sample in tqdm(enumerate(test_loader)): img, mask_2D = sample["image"], sample["valid_mask"] img = img.transpose(0, 1) img_2D = sample["image_2D"].numpy().squeeze() mask_2D = mask_2D.transpose(0, 1) inv_homographies, homographies = ( sample["homographies"], sample["inv_homographies"], ) img, mask_2D, homographies, inv_homographies = ( img.to(device), mask_2D.to(device), homographies.to(device), inv_homographies.to(device), ) # sample = test_set[i] name = sample["name"][0] logging.info(f"name: {name}") if check_exist: p = Path(save_output, "{}.npz".format(name)) if p.exists(): logging.info("file %s exists. skip the sample.", name) continue print("Pass img to network") # pass through network heatmap = fe.run(img, onlyHeatmap=True, train=False) outputs = combine_heatmap(heatmap, inv_homographies, mask_2D, device=device) pts = fe.getPtsFromHeatmap(outputs.detach().cpu().squeeze()) # (x,y, prob) # subpixel prediction if config["model"]["subpixel"]["enable"]: fe.heatmap = outputs # tensor [batch, 1, H, W] print("outputs: ", outputs.shape) print("pts: ", pts.shape) pts = fe.soft_argmax_points([pts]) pts = pts[0] ## top K points pts = pts.transpose() print("total points: ", pts.shape) print("pts: ", pts[:5]) if top_k: if pts.shape[0] > top_k: pts = pts[:top_k, :] print("topK filter: ", pts.shape) ## save keypoints pred = {} pred.update({"pts": pts}) ## - make directories filename = str(name) if task == "Kitti" or "Kitti_inh": scene_name = sample["scene_name"][0] os.makedirs(Path(save_output, scene_name), exist_ok=True) path = Path(save_output, "{}.npz".format(filename)) np.savez_compressed(path, **pred) ## output images for visualization labels if output_images: img_pts = draw_keypoints(img_2D * 255, pts.transpose()) f = save_output / (str(count) + ".png") if task == "Coco" or "Kitti": f = save_output / (name + ".png") saveImg(img_pts, str(f)) count += 1 print("output pseudo ground truth: ", count) save_file = save_output / "export.txt" with open(save_file, "a") as myfile: myfile.write("Homography adaptation: " + str(homoAdapt_iter) + "\n") myfile.write("output pairs: " + str(count) + "\n") pass
def train_joint(config, output_dir, args): assert 'train_iter' in config # config # from utils.utils import pltImshow # from utils.utils import saveImg torch.set_default_tensor_type(torch.FloatTensor) task = config['data']['dataset'] device = torch.device("cuda" if torch.cuda.is_available() else "cpu") logging.info('train on device: %s', device) with open(os.path.join(output_dir, 'config.yml'), 'w') as f: yaml.dump(config, f, default_flow_style=False) # writer = SummaryWriter(getWriterPath(task=args.command, date=True)) writer = SummaryWriter(getWriterPath(task=args.command, exper_name=args.exper_name, date=True)) ## save data save_path = get_save_path(output_dir) # let us create Dataloaders def _create_loader(dataset, bs=8, shuffle=False, n_workers=8): return torch.utils.data.DataLoader(dataset, batch_size=bs, shuffle=shuffle, pin_memory=True, num_workers=n_workers, ) splits = ["train", "val"] data_loaders = {} # create the dataset and dataloader classes for split in splits: dataset = PhototourismSuperpoint(split=split, **config["data"]) data_loaders[split] = _create_loader(dataset, config["model"]["batch_size"], split == "train", config["training"]["workers_train"]) # data loading train_loader, val_loader = data_loaders["train"], data_loaders["val"] datasize(train_loader, config, tag='train') datasize(val_loader, config, tag='val') # init the training agent using config file # from train_model_frontend import Train_model_frontend from utils.loader import get_module train_agent = TrainModelHeatmapMy(config, save_path=save_path, device=device) # writer from tensorboard train_agent.writer = writer # feed the data into the agent train_agent.train_loader = train_loader train_agent.val_loader = val_loader # load model initiates the model and load the pretrained model (if any) train_agent.loadModel() train_agent.dataParallel() try: # train function takes care of training and evaluation train_agent.train_my() except KeyboardInterrupt: print ("press ctrl + c, save model!") train_agent.saveModel() pass
def export_descriptor(config, output_dir, args): ''' 1) input 2 images, output keypoints and correspondence :param config: :param output_dir: :param args: :return: ''' # config # device = torch.device("cpu") device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") logging.info('train on device: %s', device) with open(os.path.join(output_dir, 'config.yml'), 'w') as f: yaml.dump(config, f, default_flow_style=False) writer = SummaryWriter(getWriterPath(task=args.command, date=True)) ## save data from pathlib import Path # save_path = save_path_formatter(config, output_dir) save_path = Path(output_dir) save_output = save_path save_output = save_output / 'predictions' save_path = save_path / 'checkpoints' logging.info('=> will save everything to {}'.format(save_path)) os.makedirs(save_path, exist_ok=True) os.makedirs(save_output, exist_ok=True) # data loading from utils.loader import dataLoader_test as dataLoader data = dataLoader(config, dataset='hpatches') test_set, test_loader = data['test_set'], data['test_loader'] from utils.print_tool import datasize datasize(test_loader, config, tag='test') from imageio import imread def load_as_float(path): return imread(path).astype(np.float32) / 255 def squeezeToNumpy(tensor_arr): return tensor_arr.detach().cpu().numpy().squeeze() outputMatches = True count = 0 max_length = 5 method = config['model']['method'] # tracker = PointTracker(max_length, nn_thresh=fe.nn_thresh) # for sample in tqdm(enumerate(test_loader)): for i, sample in tqdm(enumerate(test_loader)): img_0, img_1 = sample['image'], sample['warped_image'] imgs_np, imgs_fil = [], [] # first image, no matches imgs_np.append(img_0.numpy().squeeze()) imgs_np.append(img_1.numpy().squeeze()) # H, W = img.shape[1], img.shape[2] # img = img.view(1,1,H,W) ##### add opencv functions here ##### def classicalDetectors(image, method='sift'): """ # sift keyframe detectors and descriptors """ image = image*255 round_method = False if round_method == True: from models.classical_detectors_descriptors import classical_detector_descriptor # with quantization points, desc = classical_detector_descriptor(image, **{'method': method}) y, x = np.where(points) # pnts = np.stack((y, x), axis=1) pnts = np.stack((x, y), axis=1) # should be (x, y) ## collect descriptros desc = desc[y, x, :] else: # sift with subpixel accuracy from models.classical_detectors_descriptors import SIFT_det as classical_detector_descriptor pnts, desc = classical_detector_descriptor(image, image) print("desc shape: ", desc.shape) return pnts, desc pts_list = [] pts, desc_1 = classicalDetectors(imgs_np[0], method=method) pts_list.append(pts) print("total points: ", pts.shape) ''' pts: list [numpy (N, 2)] desc: list [numpy (N, 128)] ''' # save keypoints pred = {} pred.update({ # 'image_fil': imgs_fil[0], 'image': imgs_np[0], }) pred.update({'prob': pts, 'desc': desc_1}) # second image, output matches pred.update({ 'warped_image': imgs_np[1], # 'warped_image_fil': imgs_fil[1], }) pts, desc_2 = classicalDetectors(imgs_np[1], method=method) pts_list.append(pts) # if outputMatches == True: # tracker.update(pts, desc) # pred.update({'matches': matches.transpose()}) print("total points: ", pts.shape) pred.update({'warped_prob': pts, 'warped_desc': desc_2, 'homography': squeezeToNumpy(sample['homography']) }) ## get matches data = get_sift_match(sift_kps_ii=pts_list[0], sift_des_ii=desc_1, sift_kps_jj=pts_list[1], sift_des_jj=desc_2, if_BF_matcher=True) matches = data['match_quality_good'] print(f"matches: {matches.shape}") matches_all = data['match_quality_all'] pred.update({ 'matches': matches, 'matches_all': matches_all }) # clean last descriptor ''' pred: 'image': np(320,240) 'prob' (keypoints): np (N1, 2) 'desc': np (N2, 256) 'warped_image': np(320,240) 'warped_prob' (keypoints): np (N2, 2) 'warped_desc': np (N2, 256) 'homography': np (3,3) 'matches': np (N3, 4) ''' # save data from pathlib import Path filename = str(count) path = Path(save_output, '{}.npz'.format(filename)) np.savez_compressed(path, **pred) count += 1 print("output pairs: ", count) save_file = save_output / "export.txt" with open(save_file, "a") as myfile: myfile.write("output pairs: " + str(count) + '\n') pass
def simple_train(config, output_dir, args): batch_size = config['training']['batch_size_train'] epochs = config['training']['epoch_train'] learning_rate = config['training']['learning_rate'] savepath = os.path.join(output_dir, 'checkpoints') os.makedirs(savepath, exist_ok=True) with open(os.path.join(output_dir, 'config.yml'), 'w') as f: yaml.dump(config, f, default_flow_style=False) # Prepare for data loader data = dataLoader(config, dataset=config['data']['dataset'], warp_input=True) trainloader, valloader = data['train_loader'], data['val_loader'] datasize(trainloader, config, tag='train') datasize(valloader, config, tag='val') # Prepare for model model = UnSuperPoint(config) model.train() dev = torch.device("cuda") if torch.cuda.is_available() else torch.device( 'cpu') model.to(dev) # Prepare for tensorboard writer model.writer = SummaryWriter( getWriterPath(task=args.command, exper_name=args.export_name, date=True)) # Prepare for optimizer # model.optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9) model.optimizer = optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999)) whole_step = 0 total = len(trainloader) try: if config['data']['curriculum']: maxX = trainloader.dataset.config['homographies'][ 'perspective_amplitude_x'] maxY = trainloader.dataset.config['homographies'][ 'perspective_amplitude_y'] for epoch in tqdm(range(1, epochs + 1), desc='epoch'): if config['data']['curriculum']: trainloader.dataset.config['homographies'][ 'perspective_amplitude_x'] = epoch * maxX / epochs trainloader.dataset.config['homographies'][ 'perspective_amplitude_y'] = epoch * maxY / epochs tqdm.write("Max x Perspective: {:.6f}".format( trainloader.dataset.config['homographies'] ['perspective_amplitude_x'])) tqdm.write("Max y Perspective: {:.6f}".format( trainloader.dataset.config['homographies'] ['perspective_amplitude_y'])) for batch_idx, (img0, img1, mat) in tqdm(enumerate(trainloader), desc='step', total=total): whole_step += 1 model.step = whole_step loss = model.train_val_step(img0, img1, mat, 'train') tqdm.write('Loss: {:.6f}'.format(loss)) if whole_step % config['save_interval'] == 0: torch.save( model.state_dict(), os.path.join( savepath, config['model']['name'] + '_{}.pkl'.format(whole_step))) if args.eval and whole_step % config[ 'validation_interval'] == 0: for j, (img0, img1, mat) in enumerate(valloader): model.train_val_step(img0, img1, mat, 'valid') if j > config['training'].get("step_val", 4): break torch.save( model.state_dict(), os.path.join( savepath, config['model']['name'] + '_{}.pkl'.format(whole_step))) except KeyboardInterrupt: print("press ctrl + c, save model!") torch.save( model.state_dict(), os.path.join( savepath, config['model']['name'] + '_{}.pkl'.format(whole_step))) pass