[ { "path": ".gitignore", "content": "__pycache__/\nCAPS/__pycache__/\njupyter/.ipynb_checkpoints/\njupyter/__pycache__/\nlogs/\nout/\nmodels/\n" }, { "path": "CAPS/__init__.py", "content": "" }, { "path": "CAPS/caps_model.py", "content": "import os\nimport cv2\nimport numpy as np\nimport torch\nfrom torch import optim\nfrom torch.autograd import Variable\nimport utils\nimport torchvision.utils as vutils\nfrom CAPS.criterion import CtoFCriterion\nfrom CAPS.network import CAPSNet\n\n\nclass CAPSModel():\n\n def name(self):\n return 'CAPS Model'\n\n def __init__(self, args):\n self.args = args\n self.device = 'cuda' if torch.cuda.is_available() else 'cpu'\n\n # init model, optimizer, scheduler\n self.model = CAPSNet(args, self.device)\n self.optimizer = optim.Adam(self.model.parameters(), lr=self.args.lr)\n self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer,\n step_size=args.lrate_decay_steps,\n gamma=args.lrate_decay_factor)\n # reloading from checkpoints\n self.start_step = self.load_from_ckpt()\n\n # define loss function\n self.criterion = CtoFCriterion(args).to(self.device)\n\n def set_input(self, data):\n self.im1 = Variable(data['im1'].to(self.device))\n self.im2 = Variable(data['im2'].to(self.device))\n\n self.coord1 = Variable(data['coord1'].to(self.device))\n self.fmatrix = data['F'].cuda()\n self.pose = Variable(data['pose'].to(self.device))\n self.intrinsic1 = data['intrinsic1'].to(self.device)\n self.intrinsic2 = data['intrinsic2'].to(self.device)\n\n self.im1_ori = data['im1_ori']\n self.im2_ori = data['im2_ori']\n self.batch_size = len(self.im1)\n self.imsize = self.im1.size()[2:]\n\n def forward(self):\n self.out = self.model.forward(self.im1, self.im2, self.coord1)\n\n def backward_net(self):\n loss = self.criterion(self.coord1, self.out, self.fmatrix, self.pose, self.imsize)\n self.j_loss, self.eloss_c, self.eloss_f, self.closs_c, self.closs_f, self.std_loss = loss\n self.j_loss.backward()\n\n def optimize_parameters(self):\n self.optimizer.zero_grad()\n self.forward()\n self.backward_net()\n self.optimizer.step()\n self.scheduler.step()\n\n def test(self):\n self.model.eval()\n with torch.no_grad():\n coord2_e, std = self.model.test(self.im1, self.im2, self.coord1)\n return coord2_e, std\n\n def extract_features(self, im, coord):\n self.model.eval()\n with torch.no_grad():\n feat_c, feat_f = self.model.extract_features(im, coord)\n return feat_c, feat_f\n\n def write_summary(self, writer, n_iter):\n print(\"%s | Step: %d, Loss: %2.5f\" % (self.args.exp_name, n_iter, self.j_loss.item()))\n\n # write scalar\n if n_iter % self.args.log_scalar_interval == 0:\n writer.add_scalar('Total_loss', self.j_loss.item(), n_iter)\n writer.add_scalar('epipolar_loss_coarse', self.eloss_c.item(), n_iter)\n writer.add_scalar('epipolar_loss_fine', self.eloss_f.item(), n_iter)\n writer.add_scalar('cycle_loss_coarse', self.closs_c.item(), n_iter)\n writer.add_scalar('cycle_loss_fine', self.closs_f.item(), n_iter)\n\n # write image\n if n_iter % self.args.log_img_interval == 0:\n # this visualization shows a number of query points in the first image,\n # and their predicted correspondences in the second image,\n # the groundtruth epipolar lines for the query points are plotted in the second image\n num_kpts_display = 20\n im1_o = self.im1_ori[0].numpy()\n im2_o = self.im2_ori[0].numpy()\n kpt1 = self.coord1.cpu().numpy()[0][:num_kpts_display, :]\n # predicted correspondence\n correspondence = self.out['coord2_ef']\n kpt2 = correspondence.detach().cpu().numpy()[0][:num_kpts_display, :]\n lines2 = cv2.computeCorrespondEpilines(kpt1.reshape(-1, 1, 2), 1, self.fmatrix[0].cpu().numpy())\n lines2 = lines2.reshape(-1, 3)\n im2_o, im1_o = utils.drawlines(im2_o, im1_o, lines2, kpt2, kpt1)\n vis = np.concatenate((im1_o, im2_o), 1)\n vis = torch.from_numpy(vis.transpose(2, 0, 1)).float().unsqueeze(0)\n x = vutils.make_grid(vis, normalize=True)\n writer.add_image('Image', x, n_iter)\n\n def load_model(self, filename):\n to_load = torch.load(filename)\n self.model.load_state_dict(to_load['state_dict'])\n if 'optimizer' in to_load.keys():\n self.optimizer.load_state_dict(to_load['optimizer'])\n if 'scheduler' in to_load.keys():\n self.scheduler.load_state_dict(to_load['scheduler'])\n return to_load['step']\n\n def load_from_ckpt(self):\n '''\n load model from existing checkpoints and return the current step\n :param ckpt_dir: the directory that stores ckpts\n :return: the current starting step\n '''\n\n # load from the specified ckpt path\n if self.args.ckpt_path != \"\":\n print(\"Reloading from {}\".format(self.args.ckpt_path))\n if os.path.isfile(self.args.ckpt_path):\n step = self.load_model(self.args.ckpt_path)\n else:\n raise Exception('no checkpoint found in the following path:{}'.format(self.args.ckpt_path))\n\n else:\n ckpt_folder = os.path.join(self.args.outdir, self.args.exp_name)\n os.makedirs(ckpt_folder, exist_ok=True)\n # load from the most recent ckpt from all existing ckpts\n ckpts = [os.path.join(ckpt_folder, f) for f in sorted(os.listdir(ckpt_folder)) if f.endswith('.pth')]\n if len(ckpts) > 0:\n fpath = ckpts[-1]\n step = self.load_model(fpath)\n print('Reloading from {}, starting at step={}'.format(fpath, step))\n else:\n print('No ckpts found, training from scratch...')\n step = 0\n\n return step\n\n def save_model(self, step):\n ckpt_folder = os.path.join(self.args.outdir, self.args.exp_name)\n os.makedirs(ckpt_folder, exist_ok=True)\n\n save_path = os.path.join(ckpt_folder, \"{:06d}.pth\".format(step))\n print('saving ckpts {}...'.format(save_path))\n torch.save({'step': step,\n 'state_dict': self.model.state_dict(),\n 'optimizer': self.optimizer.state_dict(),\n 'scheduler': self.scheduler.state_dict(),\n },\n save_path)\n" }, { "path": "CAPS/criterion.py", "content": "import torch\nimport torch.nn as nn\n\n\nclass CtoFCriterion(nn.Module):\n def __init__(self, args):\n super(CtoFCriterion, self).__init__()\n self.args = args\n self.w_ec = args.w_epipolar_coarse\n self.w_ef = args.w_epipolar_fine\n self.w_cc = args.w_cycle_coarse\n self.w_cf = args.w_cycle_coarse\n self.w_std = args.w_std\n\n def homogenize(self, coord):\n coord = torch.cat((coord, torch.ones_like(coord[:, :, [0]])), -1)\n return coord\n\n def set_weight(self, std, mask=None, regularizer=0.0):\n if self.args.std:\n inverse_std = 1. / torch.clamp(std+regularizer, min=1e-10)\n weight = inverse_std / torch.mean(inverse_std)\n weight = weight.detach() # Bxn\n else:\n weight = torch.ones_like(std)\n\n if mask is not None:\n weight *= mask.float()\n weight /= (torch.mean(weight) + 1e-8)\n return weight\n\n def epipolar_cost(self, coord1, coord2, fmatrix):\n coord1_h = self.homogenize(coord1).transpose(1, 2)\n coord2_h = self.homogenize(coord2).transpose(1, 2)\n epipolar_line = fmatrix.bmm(coord1_h) # Bx3xn\n epipolar_line_ = epipolar_line / torch.clamp(torch.norm(epipolar_line[:, :2, :], dim=1, keepdim=True), min=1e-8)\n essential_cost = torch.abs(torch.sum(coord2_h * epipolar_line_, dim=1)) # Bxn\n return essential_cost\n\n def epipolar_loss(self, coord1, coord2, fmatrix, weight):\n essential_cost = self.epipolar_cost(coord1, coord2, fmatrix)\n loss = torch.mean(weight * essential_cost)\n return loss\n\n def cycle_consistency_loss(self, coord1, coord1_loop, weight, th=40):\n '''\n compute the cycle consistency loss\n :param coord1: [batch_size, n_pts, 2]\n :param coord1_loop: the predicted location [batch_size, n_pts, 2]\n :param weight: the weight [batch_size, n_pts]\n :param th: the threshold, only consider distances under this threshold\n :return: the cycle consistency loss value\n '''\n distance = torch.norm(coord1 - coord1_loop, dim=-1)\n distance_ = torch.zeros_like(distance)\n distance_[distance < th] = distance[distance < th]\n loss = torch.mean(weight * distance_)\n return loss\n\n def forward(self, coord1, data, fmatrix, pose, im_size):\n coord2_ec = data['coord2_ec']\n coord2_ef = data['coord2_ef']\n coord1_lc = data['coord1_lc']\n coord1_lf = data['coord1_lf']\n std_c = data['std_c']\n std_f = data['std_f']\n std_lc = data['std_lc']\n std_lf = data['std_lf']\n shorter_edge, longer_edge = min(im_size), max(im_size)\n\n epipolar_cost_c = self.epipolar_cost(coord1, coord2_ec, fmatrix)\n # only add fine level loss if the coarse level prediction is close enough to gt epipolar line\n mask_ctof = (epipolar_cost_c < (shorter_edge * self.args.window_size))\n # only add cycle consistency loss if the coarse level prediction is close enough to gt epipolar line\n mask_epip_c = (epipolar_cost_c < (shorter_edge * self.args.th_epipolar))\n mask_cycle_c = (epipolar_cost_c < (shorter_edge * self.args.th_cycle))\n\n epipolar_cost_f = self.epipolar_cost(coord1, coord2_ef, fmatrix)\n # only add cycle consistency loss if the fine level prediction is close enough to gt epipolar line\n mask_epip_f = (epipolar_cost_f < (shorter_edge * self.args.th_epipolar))\n mask_cycle_f = (epipolar_cost_f < shorter_edge * self.args.th_cycle)\n\n weight_c = self.set_weight(std_c, mask=mask_epip_c)\n weight_f = self.set_weight(std_f, mask=mask_epip_f*mask_ctof)\n\n eloss_c = torch.mean(epipolar_cost_c * weight_c) / longer_edge\n eloss_f = torch.mean(epipolar_cost_f * weight_f) / longer_edge\n\n weight_cycle_c = self.set_weight(std_c * std_lc, mask=mask_cycle_c)\n weight_cycle_f = self.set_weight(std_f * std_lf, mask=mask_cycle_f)\n\n closs_c = self.cycle_consistency_loss(coord1, coord1_lc, weight_cycle_c) / longer_edge\n closs_f = self.cycle_consistency_loss(coord1, coord1_lf, weight_cycle_f) / longer_edge\n\n loss = self.w_ec * eloss_c + self.w_ef * eloss_f + self.w_cc * closs_c + self.w_cf * closs_f\n\n std_loss = torch.mean(std_c) + torch.mean(std_f)\n loss += self.w_std * std_loss\n\n return loss, eloss_c, eloss_f, closs_c, closs_f, std_loss\n\n\n" }, { "path": "CAPS/network.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport importlib\n\n\nclass CAPSNet(nn.Module):\n def __init__(self, args, device):\n super(CAPSNet, self).__init__()\n self.args = args\n self.device = device\n self.net = ResUNet(pretrained=args.pretrained,\n encoder=args.backbone,\n coarse_out_ch=args.coarse_feat_dim,\n fine_out_ch=args.fine_feat_dim).to(self.device)\n\n @staticmethod\n def normalize(coord, h, w):\n '''\n turn the coordinates from pixel indices to the range of [-1, 1]\n :param coord: [..., 2]\n :param h: the image height\n :param w: the image width\n :return: the normalized coordinates [..., 2]\n '''\n c = torch.Tensor([(w - 1) / 2., (h - 1) / 2.]).to(coord.device).float()\n coord_norm = (coord - c) / c\n return coord_norm\n\n @staticmethod\n def denormalize(coord_norm, h, w):\n '''\n turn the coordinates from normalized value ([-1, 1]) to actual pixel indices\n :param coord_norm: [..., 2]\n :param h: the image height\n :param w: the image width\n :return: actual pixel coordinates\n '''\n c = torch.Tensor([(w - 1) / 2., (h - 1) / 2.]).to(coord_norm.device)\n coord = coord_norm * c + c\n return coord\n\n def ind2coord(self, ind, width):\n ind = ind.unsqueeze(-1)\n x = ind % width\n y = ind // width\n coord = torch.cat((x, y), -1).float()\n return coord\n\n def gen_grid(self, h_min, h_max, w_min, w_max, len_h, len_w):\n x, y = torch.meshgrid([torch.linspace(w_min, w_max, len_w), torch.linspace(h_min, h_max, len_h)])\n grid = torch.stack((x, y), -1).transpose(0, 1).reshape(-1, 2).float().to(self.device)\n return grid\n\n def sample_feat_by_coord(self, x, coord_n, norm=False):\n '''\n sample from normalized coordinates\n :param x: feature map [batch_size, n_dim, h, w]\n :param coord_n: normalized coordinates, [batch_size, n_pts, 2]\n :param norm: if l2 normalize features\n :return: the extracted features, [batch_size, n_pts, n_dim]\n '''\n feat = F.grid_sample(x, coord_n.unsqueeze(2)).squeeze(-1)\n if norm:\n feat = F.normalize(feat)\n feat = feat.transpose(1, 2)\n return feat\n\n def compute_prob(self, feat1, feat2):\n '''\n compute probability\n :param feat1: query features, [batch_size, m, n_dim]\n :param feat2: reference features, [batch_size, n, n_dim]\n :return: probability, [batch_size, m, n]\n '''\n assert self.args.prob_from in ['correlation', 'distance']\n if self.args.prob_from == 'correlation':\n sim = feat1.bmm(feat2.transpose(1, 2))\n prob = F.softmax(sim, dim=-1) # Bxmxn\n else:\n dist = torch.sum(feat1**2, dim=-1, keepdim=True) + \\\n torch.sum(feat2**2, dim=-1, keepdim=True).transpose(1, 2) - \\\n 2 * feat1.bmm(feat2.transpose(1, 2))\n prob = F.softmax(-dist, dim=-1) # Bxmxn\n return prob\n\n def get_1nn_coord(self, feat1, featmap2):\n '''\n find the coordinates of nearest neighbor match\n :param feat1: query features, [batch_size, n_pts, n_dim]\n :param featmap2: the feature maps of the other image\n :return: normalized correspondence locations [batch_size, n_pts, 2]\n '''\n batch_size, d, h, w = featmap2.shape\n feat2_flatten = featmap2.reshape(batch_size, d, h*w).transpose(1, 2) # Bx(hw)xd\n\n assert self.args.prob_from in ['correlation', 'distance']\n if self.args.prob_from == 'correlation':\n sim = feat1.bmm(feat2_flatten.transpose(1, 2))\n ind2_1nn = torch.max(sim, dim=-1)[1]\n else:\n dist = torch.sum(feat1**2, dim=-1, keepdim=True) + \\\n torch.sum(feat2_flatten**2, dim=-1, keepdim=True).transpose(1, 2) - \\\n 2 * feat1.bmm(feat2_flatten.transpose(1, 2))\n ind2_1nn = torch.min(dist, dim=-1)[1]\n\n coord2 = self.ind2coord(ind2_1nn, w)\n coord2_n = self.normalize(coord2, h, w)\n return coord2_n\n\n def get_expected_correspondence_locs(self, feat1, featmap2, with_std=False):\n '''\n compute the expected correspondence locations\n :param feat1: the feature vectors of query points [batch_size, n_pts, n_dim]\n :param featmap2: the feature maps of the reference image [batch_size, n_dim, h, w]\n :param with_std: if return the standard deviation\n :return: the normalized expected correspondence locations [batch_size, n_pts, 2]\n '''\n B, d, h2, w2 = featmap2.size()\n grid_n = self.gen_grid(-1, 1, -1, 1, h2, w2)\n featmap2_flatten = featmap2.reshape(B, d, h2*w2).transpose(1, 2) # BX(hw)xd\n prob = self.compute_prob(feat1, featmap2_flatten) # Bxnx(hw)\n\n grid_n = grid_n.unsqueeze(0).unsqueeze(0) # 1x1x(hw)x2\n expected_coord_n = torch.sum(grid_n * prob.unsqueeze(-1), dim=2) # Bxnx2\n\n if with_std:\n # convert to normalized scale [-1, 1]\n var = torch.sum(grid_n**2 * prob.unsqueeze(-1), dim=2) - expected_coord_n**2 # Bxnx2\n std = torch.sum(torch.sqrt(torch.clamp(var, min=1e-10)), -1) # Bxn\n return expected_coord_n, std\n else:\n return expected_coord_n\n\n def get_expected_correspondence_within_window(self, feat1, featmap2, coord2_n, with_std=False):\n '''\n :param feat1: the feature vectors of query points [batch_size, n_pts, n_dim]\n :param featmap2: the feature maps of the reference image [batch_size, n_dim, h, w]\n :param coord2_n: normalized center locations [batch_size, n_pts, 2]\n :param with_std: if return the standard deviation\n :return: the normalized expected correspondence locations, [batch_size, n_pts, 2], optionally with std\n '''\n batch_size, n_dim, h2, w2 = featmap2.shape\n n_pts = coord2_n.shape[1]\n grid_n = self.gen_grid(h_min=-self.args.window_size, h_max=self.args.window_size,\n w_min=-self.args.window_size, w_max=self.args.window_size,\n len_h=int(self.args.window_size*h2), len_w=int(self.args.window_size*w2))\n\n grid_n_ = grid_n.repeat(batch_size, 1, 1, 1) # Bx1xhwx2\n coord2_n_grid = coord2_n.unsqueeze(-2) + grid_n_ # Bxnxhwx2\n feat2_win = F.grid_sample(featmap2, coord2_n_grid, padding_mode='zeros').permute(0, 2, 3, 1) # Bxnxhwxd\n\n feat1 = feat1.unsqueeze(-2)\n\n prob = self.compute_prob(feat1.reshape(batch_size*n_pts, -1, n_dim),\n feat2_win.reshape(batch_size*n_pts, -1, n_dim)).reshape(batch_size, n_pts, -1)\n\n expected_coord2_n = torch.sum(coord2_n_grid * prob.unsqueeze(-1), dim=2) # Bxnx2\n\n if with_std:\n var = torch.sum(coord2_n_grid**2 * prob.unsqueeze(-1), dim=2) - expected_coord2_n**2 # Bxnx2\n std = torch.sum(torch.sqrt(torch.clamp(var, min=1e-10)), -1) # Bxn\n return expected_coord2_n, std\n else:\n return expected_coord2_n\n\n def forward(self, im1, im2, coord1):\n # extract features for both images\n xc1, xf1 = self.net(im1)\n xc2, xf2 = self.net(im2)\n\n # image width and height\n h1i, w1i = im1.size()[2:]\n h2i, w2i = im2.size()[2:]\n\n coord1_n = self.normalize(coord1, h1i, w1i)\n feat1_coarse = self.sample_feat_by_coord(xc1, coord1_n) # Bxnxd\n coord2_ec_n, std_c = self.get_expected_correspondence_locs(feat1_coarse, xc2, with_std=True)\n\n # the center locations of the local window for fine level computation\n coord2_ec_n_ = self.get_1nn_coord(feat1_coarse, xc2) if self.args.use_nn else coord2_ec_n\n feat1_fine = self.sample_feat_by_coord(xf1, coord1_n) # Bxnxd\n coord2_ef_n, std_f = self.get_expected_correspondence_within_window(feat1_fine, xf2,\n coord2_ec_n_, with_std=True)\n\n feat2_coarse = self.sample_feat_by_coord(xc2, coord2_ec_n_)\n coord1_lc_n, std_lc = self.get_expected_correspondence_locs(feat2_coarse, xc1, with_std=True)\n\n feat2_fine = self.sample_feat_by_coord(xf2, coord2_ef_n) # Bxnxd\n coord1_lf_n, std_lf = self.get_expected_correspondence_within_window(feat2_fine, xf1,\n coord1_n, with_std=True)\n coord2_ec = self.denormalize(coord2_ec_n, h2i, w2i)\n coord2_ef = self.denormalize(coord2_ef_n, h2i, w2i)\n coord1_lc = self.denormalize(coord1_lc_n, h1i, w1i)\n coord1_lf = self.denormalize(coord1_lf_n, h1i, w1i)\n\n return {'coord2_ec': coord2_ec, 'coord2_ef': coord2_ef,\n 'coord1_lc': coord1_lc, 'coord1_lf': coord1_lf,\n 'std_c': std_c, 'std_f': std_f,\n 'std_lc': std_lc, 'std_lf': std_lf,\n }\n\n def extract_features(self, im, coord):\n '''\n extract coarse and fine level features given the input image and 2d locations\n :param im: [batch_size, 3, h, w]\n :param coord: [batch_size, n_pts, 2]\n :return: coarse features [batch_size, n_pts, coarse_feat_dim] and fine features [batch_size, n_pts, fine_feat_dim]\n '''\n xc, xf = self.net(im)\n hi, wi = im.size()[2:]\n coord_n = self.normalize(coord, hi, wi)\n feat_c = self.sample_feat_by_coord(xc, coord_n)\n feat_f = self.sample_feat_by_coord(xf, coord_n)\n return feat_c, feat_f\n\n def test(self, im1, im2, coord1):\n '''\n given a pair of images im1, im2, compute the coorrespondences for query points coord1.\n We performa full image search at coarse level and local search at fine level\n :param im1: [batch_size, 3, h, w]\n :param im2: [batch_size, 3, h, w]\n :param coord1: [batch_size, n_pts, 2]\n :return: the fine level correspondence location [batch_size, n_pts, 2]\n '''\n\n xc1, xf1 = self.net(im1)\n xc2, xf2 = self.net(im2)\n\n h1i, w1i = im1.shape[2:]\n h2i, w2i = im2.shape[2:]\n\n coord1_n = self.normalize(coord1, h1i, w1i)\n feat1_c = self.sample_feat_by_coord(xc1, coord1_n)\n _, std_c = self.get_expected_correspondence_locs(feat1_c, xc2, with_std=True)\n\n coord2_ec_n = self.get_1nn_coord(feat1_c, xc2)\n feat1_f = self.sample_feat_by_coord(xf1, coord1_n)\n _, std_f = self.get_expected_correspondence_within_window(feat1_f, xf2, coord2_ec_n, with_std=True)\n\n coord2_ef_n = self.get_1nn_coord(feat1_f, xf2)\n coord2_ef = self.denormalize(coord2_ef_n, h2i, w2i)\n std = (std_c + std_f)/2\n\n return coord2_ef, std\n\n\n####################### ResUnet ##########################\n\ndef class_for_name(module_name, class_name):\n # load the module, will raise ImportError if module cannot be loaded\n m = importlib.import_module(module_name)\n return getattr(m, class_name)\n\n\nclass conv(nn.Module):\n def __init__(self, num_in_layers, num_out_layers, kernel_size, stride):\n super(conv, self).__init__()\n self.kernel_size = kernel_size\n self.conv = nn.Conv2d(num_in_layers,\n num_out_layers,\n kernel_size=kernel_size,\n stride=stride,\n padding=(self.kernel_size - 1) // 2)\n self.bn = nn.BatchNorm2d(num_out_layers)\n\n def forward(self, x):\n return F.elu(self.bn(self.conv(x)), inplace=True)\n\n\nclass upconv(nn.Module):\n def __init__(self, num_in_layers, num_out_layers, kernel_size, scale):\n super(upconv, self).__init__()\n self.scale = scale\n self.conv = conv(num_in_layers, num_out_layers, kernel_size, 1)\n\n def forward(self, x):\n x = nn.functional.interpolate(x, scale_factor=self.scale, align_corners=True, mode='bilinear')\n return self.conv(x)\n\n\nclass ResUNet(nn.Module):\n def __init__(self,\n encoder='resnet50',\n pretrained=True,\n coarse_out_ch=128,\n fine_out_ch=128\n ):\n\n super(ResUNet, self).__init__()\n assert encoder in ['resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152'], \"Incorrect encoder type\"\n if encoder in ['resnet18', 'resnet34']:\n filters = [64, 128, 256, 512]\n else:\n filters = [256, 512, 1024, 2048]\n resnet = class_for_name(\"torchvision.models\", encoder)(pretrained=pretrained)\n\n self.firstconv = resnet.conv1 # H/2\n self.firstbn = resnet.bn1\n self.firstrelu = resnet.relu\n self.firstmaxpool = resnet.maxpool # H/4\n\n # encoder\n self.layer1 = resnet.layer1 # H/4\n self.layer2 = resnet.layer2 # H/8\n self.layer3 = resnet.layer3 # H/16\n\n # coarse-level conv\n self.conv_coarse = conv(filters[2], coarse_out_ch, 1, 1)\n\n # decoder\n self.upconv3 = upconv(filters[2], 512, 3, 2)\n self.iconv3 = conv(filters[1] + 512, 512, 3, 1)\n self.upconv2 = upconv(512, 256, 3, 2)\n self.iconv2 = conv(filters[0] + 256, 256, 3, 1)\n\n # fine-level conv\n self.conv_fine = conv(256, fine_out_ch, 1, 1)\n\n def skipconnect(self, x1, x2):\n diffY = x2.size()[2] - x1.size()[2]\n diffX = x2.size()[3] - x1.size()[3]\n\n x1 = F.pad(x1, (diffX // 2, diffX - diffX // 2,\n diffY // 2, diffY - diffY // 2))\n\n # for padding issues, see\n # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a\n # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd\n\n x = torch.cat([x2, x1], dim=1)\n return x\n\n def forward(self, x):\n x = self.firstrelu(self.firstbn(self.firstconv(x)))\n x = self.firstmaxpool(x)\n\n x1 = self.layer1(x)\n x2 = self.layer2(x1)\n x3 = self.layer3(x2)\n\n x_coarse = self.conv_coarse(x3)\n\n x = self.upconv3(x3)\n x = self.skipconnect(x2, x)\n x = self.iconv3(x)\n\n x = self.upconv2(x)\n x = self.skipconnect(x1, x)\n x = self.iconv2(x)\n\n x_fine = self.conv_fine(x)\n return [x_coarse, x_fine]\n" }, { "path": "LICENSE", "content": "Copyright 2021 Qianqian Wang\n\nPermission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the \"Software\"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.\n" }, { "path": "README.md", "content": "# Learning Feature Descriptors using Camera Pose Supervision\n\nThis repository contains a PyTorch implementation of the paper:\n\n[*Learning Feature Descriptors using Camera Pose Supervision*](https://qianqianwang68.github.io/CAPS/)\n[[Project page]](https://qianqianwang68.github.io/CAPS/)\n[[Arxiv]](https://arxiv.org/abs/2004.13324)\n\n[Qianqian Wang](https://www.cs.cornell.edu/~qqw/), \n[Xiaowei Zhou](http://www.cad.zju.edu.cn/home/xzhou/),\n[Bharath Hariharan](http://home.bharathh.info/),\n[Noah Snavely](http://www.cs.cornell.edu/~snavely/) \n\nECCV 2020 (*Oral*)\n\n![Teaser](assets/teaser.jpg)\n\n## Abstract\n\nRecent research on learned visual descriptors has shown promising improvements in correspondence estimation, a key component of many 3D vision tasks. However, existing descriptor learning frameworks typically require ground-truth correspondences between feature points for training, which are challenging to acquire at scale. In this paper we propose a novel weakly-supervised framework that can learn feature descriptors solely from relative camera poses between images. To do so, we devise both a new loss function that exploits the epipolar constraint given by camera poses, and a new model architecture that makes the whole pipeline differentiable and efficient. Because we no longer need pixel-level ground-truth correspondences, our framework opens up the possibility of training on much larger and more diverse datasets for better and unbiased descriptors. We call the resulting descriptors CAmera Pose Supervised, or CAPS, descriptors. Though trained with weak supervision, CAPS descriptors outperform even prior fully-supervised descriptors and achieve state-of-the-art performance on a variety of geometric tasks.\n\n\n## Requirements\n```bash\n# Create conda environment with torch 1.0.1 and CUDA 10.0\nconda env create -f environment.yml\nconda activate caps\n```\nIf you encounter problems with OpenCV, try to uninstall your current opencv packages and reinstall them again\n```bash\npip uninstall opencv-python\npip uninstall opencv-contrib-python\npip install opencv-python==3.4.2.17\npip install opencv-contrib-python==3.4.2.17\n```\n\n## Pretrained Model\nPretrained model can be downloaded using this google drive [link](https://drive.google.com/file/d/1KfIQYyM7vlSvvQ3X7xi1YSPrietqeXYd/view?usp=sharing), or this BaiduYun [link](https://pan.baidu.com/s/1rGt4okK3KsumPLUVhOJdXQ) (password: 3na7).\n\n\n## Dataset\nPlease download the preprocessed MegaDepth dataset using this google drive [link](https://drive.google.com/file/d/1Xp6BRKx_5sIwdJWVK0kJTmtUcQNMY5yL/view?usp=sharing) or this BaiduYun [link](https://pan.baidu.com/s/1rGt4okK3KsumPLUVhOJdXQ) (password: 3na7\n\n## Training\nTo start training, please download our [training data](https://drive.google.com/file/d/1Xp6BRKx_5sIwdJWVK0kJTmtUcQNMY5yL/view?usp=sharing), and run the following command: \n\n```bash\n# example usage\npython train.py --config configs/train_megadepth.yaml\n```\n\n\n## Feature extraction\nWe provide code for extracting CAPS descriptors on HPatches dataset.\nTo download and use the HPatches Sequences, please refer to this [link](https://github.com/mihaidusmanu/d2-net/tree/master/hpatches_sequences).\n\nTo extract CAPS features on HPatches dataset, download the pretrained model, modify paths in ```configs/extract_features_hpatches.yaml``` and run\n```bash\npython extract_features.py --config configs/extract_features_hpatches.yaml\n```\n\n## Interactive demo\nWe provide an interactive demo where you could click on locations in the first image and see their predicted correspondences in the second image.\nPlease refer to ```jupyter/visualization.ipynb``` for more details.\n\n## Cite\nPlease cite our work if you find it useful: \n```bibtex\n@inproceedings{wang2020learning,\n Title = {Learning Feature Descriptors using Camera Pose Supervision},\n Author = {Qianqian Wang and Xiaowei Zhou and Bharath Hariharan and Noah Snavely},\n booktitle = {Proc. European Conference on Computer Vision (ECCV)},\n Year = {2020},\n}\n```\n**Acknowledgements**. We thank Kai Zhang, Zixin Luo, Zhengqi Li for helpful discussion and comments. This work was partly supported by a DARPA\nLwLL grant, and in part by the generosity of Eric and Wendy Schmidt by\nrecommendation of the Schmidt Futures program.\n" }, { "path": "config.py", "content": "import configargparse\n\n\ndef get_args():\n parser = configargparse.ArgParser(config_file_parser_class=configargparse.YAMLConfigFileParser)\n parser.add_argument('--config', is_config_file=True, help='config file path')\n\n ## path options\n parser.add_argument('--datadir', type=str, help='the dataset directory')\n parser.add_argument(\"--logdir\", type=str, default='./logs/', help='dir of tensorboard logs')\n parser.add_argument(\"--outdir\", type=str, default='./out/', help='dir of output e.g., ckpts')\n parser.add_argument(\"--ckpt_path\", type=str, default=\"\",\n help='specific checkpoint path to load the model from, '\n 'if not specified, automatically reload from most recent checkpoints')\n\n ## general options\n parser.add_argument(\"--exp_name\", type=str, help='experiment name')\n parser.add_argument('--n_iters', type=int, default=200000, help='max number of training iterations')\n parser.add_argument('--phase', type=str, default='train', help='train/val/test')\n\n # data options\n parser.add_argument('--workers', type=int, help='number of data loading workers', default=4)\n parser.add_argument('--num_pts', type=int, default=500, help='num of points trained in each pair')\n parser.add_argument('--train_kp', type=str, default='mixed', help='sift/random/mixed')\n parser.add_argument('--prune_kp', type=int, default=1, help='if prune non-matchable keypoints')\n\n # training options\n parser.add_argument('--batch_size', type=int, default=6, help='input batch size')\n parser.add_argument('--lr', type=float, default=1e-4, help='base learning rate')\n parser.add_argument(\"--lrate_decay_steps\", type=int, default=80000,\n help='decay learning rate by a factor every specified number of steps')\n parser.add_argument(\"--lrate_decay_factor\", type=float, default=0.5,\n help='decay learning rate by a factor every specified number of steps')\n\n ## model options\n parser.add_argument('--backbone', type=str, default='resnet50',\n help='backbone for feature representation extraction. supported: resent')\n parser.add_argument('--pretrained', type=int, default=1,\n help='if use ImageNet pretrained weights to initialize the network')\n parser.add_argument('--coarse_feat_dim', type=int, default=128,\n help='the feature dimension for coarse level features')\n parser.add_argument('--fine_feat_dim', type=int, default=128,\n help='the feature dimension for fine level features')\n parser.add_argument('--prob_from', type=str, default='correlation',\n help='compute prob by softmax(correlation score), or softmax(-distance),'\n 'options: correlation|distance')\n parser.add_argument('--window_size', type=float, default=0.125,\n help='the size of the window, w.r.t image width at the fine level')\n parser.add_argument('--use_nn', type=int, default=1, help='if use nearest neighbor in the coarse level')\n\n ## loss function options\n parser.add_argument('--std', type=int, default=1, help='reweight loss using the standard deviation')\n parser.add_argument('--w_epipolar_coarse', type=float, default=1, help='coarse level epipolar loss weight')\n parser.add_argument('--w_epipolar_fine', type=float, default=1, help='fine level epipolar loss weight')\n parser.add_argument('--w_cycle_coarse', type=float, default=0.1, help='coarse level cycle consistency loss weight')\n parser.add_argument('--w_cycle_fine', type=float, default=0.1, help='fine level cycle consistency loss weight')\n parser.add_argument('--w_std', type=float, default=0, help='the weight for the loss on std')\n parser.add_argument('--th_cycle', type=float, default=0.025,\n help='if the distance (normalized scale) from the prediction to epipolar line > this th, '\n 'do not add the cycle consistency loss')\n parser.add_argument('--th_epipolar', type=float, default=0.5,\n help='if the distance (normalized scale) from the prediction to epipolar line > this th, '\n 'do not add the epipolar loss')\n\n ## logging options\n parser.add_argument('--log_scalar_interval', type=int, default=20, help='print interval')\n parser.add_argument('--log_img_interval', type=int, default=500, help='log image interval')\n parser.add_argument(\"--save_interval\", type=int, default=10000, help='frequency of weight ckpt saving')\n\n ## eval options\n parser.add_argument('--extract_img_dir', type=str, help='the directory of images to extract features')\n parser.add_argument('--extract_out_dir', type=str, help='the directory of images to extract features')\n\n args = parser.parse_known_args()[0]\n\n return args\n" }, { "path": "configs/extract_features_hpatches.yaml", "content": "# extract feature descriptors using pretrained model weights\n\nexpname: feature_extraction\n# replace the following with your pretrained model path, img_dir and out_dir, respectively\nckpt_path: 'pretrained/caps-pretrained.pth'\nextract_img_dir: '/phoenix/S7/qw246/hpatches-benchmark/data/d2net/hpatches-sequences-release/'\nextract_out_dir: '/phoenix/S7/qw246/hpatches-benchmark/data/output/desc/caps/'" }, { "path": "configs/train_megadepth.yaml", "content": "# training from scratch using a single gpu, default configs in config.py file\nexp_name: train_caps\ndatadir: /phoenix/S7/qw246/CAPS-MegaDepth-release-light # replace with your data dir\n" }, { "path": "dataloader/__init__.py", "content": "" }, { "path": "dataloader/data_utils.py", "content": "import numpy as np\nimport cv2\n\n\ndef skew(x):\n return np.array([[0, -x[2], x[1]],\n [x[2], 0, -x[0]],\n [-x[1], x[0], 0]])\n\n\ndef rotateImage(image, angle):\n h, w = image.shape[:2]\n angle_radius = np.abs(angle / 180. * np.pi)\n cos = np.cos(angle_radius)\n sin = np.sin(angle_radius)\n tan = np.tan(angle_radius)\n scale_h = (h / cos + (w - h * tan) * sin) / h\n scale_w = (h / sin + (w - h / tan) * cos) / w\n scale = max(scale_h, scale_w)\n image_center = tuple(np.array(image.shape[1::-1]) / 2.)\n rot_mat = cv2.getRotationMatrix2D(image_center, angle, scale)\n result = cv2.warpAffine(image, rot_mat, image.shape[1::-1], flags=cv2.INTER_LINEAR)\n rotation = np.eye(4)\n rotation[:2, :2] = rot_mat[:2, :2]\n return result, rotation\n\n\ndef perspective_transform(img, param=0.001):\n h, w = img.shape[:2]\n random_state = np.random.RandomState(None)\n M = np.array([[1 - param + 2 * param * random_state.rand(),\n -param + 2 * param * random_state.rand(),\n -param + 2 * param * random_state.rand()],\n [-param + 2 * param * random_state.rand(),\n 1 - param + 2 * param * random_state.rand(),\n -param + 2 * param * random_state.rand()],\n [-param + 2 * param * random_state.rand(),\n -param + 2 * param * random_state.rand(),\n 1 - param + 2 * param * random_state.rand()]])\n\n dst = cv2.warpPerspective(img, M, (w, h))\n return dst, M\n\n\ndef generate_query_kpts(img, mode, num_pts, h, w):\n # generate candidate query points\n if mode == 'random':\n kp1_x = np.random.rand(num_pts) * (w - 1)\n kp1_y = np.random.rand(num_pts) * (h - 1)\n coord = np.stack((kp1_x, kp1_y)).T\n\n elif mode == 'sift':\n gray1 = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)\n sift = cv2.xfeatures2d.SIFT_create(nfeatures=num_pts)\n kp1 = sift.detect(gray1)\n coord = np.array([[kp.pt[0], kp.pt[1]] for kp in kp1])\n\n elif mode == 'mixed':\n kp1_x = np.random.rand(1 * int(0.1 * num_pts)) * (w - 1)\n kp1_y = np.random.rand(1 * int(0.1 * num_pts)) * (h - 1)\n kp1_rand = np.stack((kp1_x, kp1_y)).T\n\n sift = cv2.xfeatures2d.SIFT_create(nfeatures=int(0.9 * num_pts))\n gray1 = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)\n kp1_sift = sift.detect(gray1)\n kp1_sift = np.array([[kp.pt[0], kp.pt[1]] for kp in kp1_sift])\n if len(kp1_sift) == 0:\n coord = kp1_rand\n else:\n coord = np.concatenate((kp1_rand, kp1_sift), 0)\n\n else:\n raise Exception('unknown type of keypoints')\n\n return coord\n\n\ndef prune_kpts(coord1, F_gt, im2_size, intrinsic1, intrinsic2, pose, d_min, d_max):\n # compute the epipolar lines corresponding to coord1\n coord1_h = np.concatenate([coord1, np.ones_like(coord1[:, [0]])], axis=1).T # 3xn\n epipolar_line = F_gt.dot(coord1_h) # 3xn\n epipolar_line /= np.clip(np.linalg.norm(epipolar_line[:2], axis=0), a_min=1e-10, a_max=None) # 3xn\n\n # determine whether the epipolar lines intersect with the second image\n h2, w2 = im2_size\n corners = np.array([[0, 0, 1], [0, h2 - 1, 1], [w2 - 1, 0, 1], [w2 - 1, h2 - 1, 1]]) # 4x3\n dists = np.abs(corners.dot(epipolar_line))\n # if the epipolar line is far away from any image corners than sqrt(h^2+w^2)\n # it doesn't intersect with the image\n non_intersect = (dists > np.sqrt(w2 ** 2 + h2 ** 2)).any(axis=0)\n\n # determine if points in coord1 is likely to have correspondence in the other image by the rough depth range\n intrinsic1_4x4 = np.eye(4)\n intrinsic1_4x4[:3, :3] = intrinsic1\n intrinsic2_4x4 = np.eye(4)\n intrinsic2_4x4[:3, :3] = intrinsic2\n coord1_h_min = np.concatenate([d_min * coord1,\n d_min * np.ones_like(coord1[:, [0]]),\n np.ones_like(coord1[:, [0]])], axis=1).T\n coord1_h_max = np.concatenate([d_max * coord1,\n d_max * np.ones_like(coord1[:, [0]]),\n np.ones_like(coord1[:, [0]])], axis=1).T\n coord2_h_min = intrinsic2_4x4.dot(pose).dot(np.linalg.inv(intrinsic1_4x4)).dot(coord1_h_min)\n coord2_h_max = intrinsic2_4x4.dot(pose).dot(np.linalg.inv(intrinsic1_4x4)).dot(coord1_h_max)\n coord2_min = coord2_h_min[:2] / (coord1_h_min[2] + 1e-10)\n coord2_max = coord2_h_max[:2] / (coord1_h_max[2] + 1e-10)\n out_range = ((coord2_min[0] < 0) & (coord2_max[0] < 0)) | \\\n ((coord2_min[1] < 0) & (coord2_max[1] < 0)) | \\\n ((coord2_min[0] > w2 - 1) & (coord2_max[0] > w2 - 1)) | \\\n ((coord2_min[1] > h2 - 1) & (coord2_max[1] > h2 - 1))\n\n ind_intersect = ~(non_intersect | out_range)\n return ind_intersect" }, { "path": "dataloader/megadepth.py", "content": "import torch\nfrom torch.utils.data import Dataset\nimport os\nimport numpy as np\nimport cv2\nimport skimage.io as io\nimport torchvision.transforms as transforms\nimport utils\nimport collections\nfrom tqdm import tqdm\nimport dataloader.data_utils as data_utils\n\n\nrand = np.random.RandomState(234)\n\n\nclass MegaDepthLoader():\n def __init__(self, args):\n self.args = args\n self.dataset = MegaDepth(args)\n self.data_loader = torch.utils.data.DataLoader(self.dataset, batch_size=args.batch_size, shuffle=False,\n num_workers=args.workers, collate_fn=self.my_collate)\n\n def my_collate(self, batch):\n ''' Puts each data field into a tensor with outer dimension batch size '''\n batch = list(filter(lambda b: b is not None, batch))\n return torch.utils.data.dataloader.default_collate(batch)\n\n def load_data(self):\n return self.data_loader\n\n def name(self):\n return 'MegaDepthLoader'\n\n def __len__(self):\n return len(self.dataset)\n\n\nclass MegaDepth(Dataset):\n def __init__(self, args):\n self.args = args\n if args.phase == 'train':\n # augment during training\n self.transform = transforms.Compose([transforms.ToPILImage(),\n transforms.ColorJitter\n (brightness=1, contrast=1, saturation=1, hue=0.4),\n transforms.ToTensor(),\n transforms.Normalize(mean=(0.485, 0.456, 0.406),\n std=(0.229, 0.224, 0.225)),\n ])\n else:\n self.transform = transforms.Compose([transforms.ToTensor(),\n transforms.Normalize(mean=(0.485, 0.456, 0.406),\n std=(0.229, 0.224, 0.225)),\n ])\n self.phase = args.phase\n self.root = os.path.join(args.datadir, self.phase)\n self.images = self.read_img_cam()\n self.imf1s, self.imf2s = self.read_pairs()\n print('total number of image pairs loaded: {}'.format(len(self.imf1s)))\n # shuffle data\n index = np.arange(len(self.imf1s))\n rand.shuffle(index)\n self.imf1s = list(np.array(self.imf1s)[index])\n self.imf2s = list(np.array(self.imf2s)[index])\n\n def read_img_cam(self):\n images = {}\n Image = collections.namedtuple(\n \"Image\", [\"name\", \"w\", \"h\", \"fx\", \"fy\", \"cx\", \"cy\", \"rvec\", \"tvec\"])\n for scene_id in os.listdir(self.root):\n densefs = [f for f in os.listdir(os.path.join(self.root, scene_id))\n if 'dense' in f and os.path.isdir(os.path.join(self.root, scene_id, f))]\n for densef in densefs:\n folder = os.path.join(self.root, scene_id, densef, 'aligned')\n img_cam_txt_path = os.path.join(folder, 'img_cam.txt')\n with open(img_cam_txt_path, \"r\") as fid:\n while True:\n line = fid.readline()\n if not line:\n break\n line = line.strip()\n if len(line) > 0 and line[0] != \"#\":\n elems = line.split()\n image_name = elems[0]\n img_path = os.path.join(folder, 'images', image_name)\n w, h = int(elems[1]), int(elems[2])\n fx, fy = float(elems[3]), float(elems[4])\n cx, cy = float(elems[5]), float(elems[6])\n R = np.array(elems[7:16])\n T = np.array(elems[16:19])\n images[img_path] = Image(\n name=image_name, w=w, h=h, fx=fx, fy=fy, cx=cx, cy=cy, rvec=R, tvec=T\n )\n return images\n\n def read_pairs(self):\n imf1s, imf2s = [], []\n print('reading image pairs from {}...'.format(self.root))\n for scene_id in tqdm(os.listdir(self.root), desc='# loading data from scene folders'):\n densefs = [f for f in os.listdir(os.path.join(self.root, scene_id))\n if 'dense' in f and os.path.isdir(os.path.join(self.root, scene_id, f))]\n for densef in densefs:\n imf1s_ = []\n imf2s_ = []\n folder = os.path.join(self.root, scene_id, densef, 'aligned')\n pairf = os.path.join(folder, 'pairs.txt')\n\n if os.path.exists(pairf):\n f = open(pairf, 'r')\n for line in f:\n imf1, imf2 = line.strip().split(' ')\n imf1s_.append(os.path.join(folder, 'images', imf1))\n imf2s_.append(os.path.join(folder, 'images', imf2))\n\n # make # image pairs per scene more balanced\n if len(imf1s_) > 5000:\n index = np.arange(len(imf1s_))\n rand.shuffle(index)\n imf1s_ = list(np.array(imf1s_)[index[:5000]])\n imf2s_ = list(np.array(imf2s_)[index[:5000]])\n\n imf1s.extend(imf1s_)\n imf2s.extend(imf2s_)\n\n return imf1s, imf2s\n\n @staticmethod\n def get_intrinsics(im_meta):\n return np.array([[im_meta.fx, 0, im_meta.cx],\n [0, im_meta.fy, im_meta.cy],\n [0, 0, 1]])\n\n @staticmethod\n def get_extrinsics(im_meta):\n R = im_meta.rvec.reshape(3, 3)\n t = im_meta.tvec\n extrinsic = np.eye(4)\n extrinsic[:3, :3] = R\n extrinsic[:3, 3] = t\n return extrinsic\n\n def __getitem__(self, item):\n imf1 = self.imf1s[item]\n imf2 = self.imf2s[item]\n im1_meta = self.images[imf1]\n im2_meta = self.images[imf2]\n im1 = io.imread(imf1)\n im2 = io.imread(imf2)\n h, w = im1.shape[:2]\n\n intrinsic1 = self.get_intrinsics(im1_meta)\n intrinsic2 = self.get_intrinsics(im2_meta)\n\n extrinsic1 = self.get_extrinsics(im1_meta)\n extrinsic2 = self.get_extrinsics(im2_meta)\n\n relative = extrinsic2.dot(np.linalg.inv(extrinsic1))\n R = relative[:3, :3]\n # remove pairs that have a relative rotation angle larger than 80 degrees\n theta = np.arccos(np.clip((np.trace(R) - 1) / 2, -1, 1)) * 180 / np.pi\n if theta > 80 and self.phase == 'train':\n return None\n\n T = relative[:3, 3]\n tx = data_utils.skew(T)\n E_gt = np.dot(tx, R)\n F_gt = np.linalg.inv(intrinsic2).T.dot(E_gt).dot(np.linalg.inv(intrinsic1))\n\n # generate candidate query points\n coord1 = data_utils.generate_query_kpts(im1, self.args.train_kp, 10*self.args.num_pts, h, w)\n\n # if no keypoints are detected\n if len(coord1) == 0:\n return None\n\n # prune query keypoints that are not likely to have correspondence in the other image\n if self.args.prune_kp:\n ind_intersect = data_utils.prune_kpts(coord1, F_gt, im2.shape[:2], intrinsic1, intrinsic2,\n relative, d_min=4, d_max=400)\n if np.sum(ind_intersect) == 0:\n return None\n coord1 = coord1[ind_intersect]\n\n coord1 = utils.random_choice(coord1, self.args.num_pts)\n coord1 = torch.from_numpy(coord1).float()\n\n im1_ori, im2_ori = torch.from_numpy(im1), torch.from_numpy(im2)\n\n F_gt = torch.from_numpy(F_gt).float() / (F_gt[-1, -1] + 1e-10)\n intrinsic1 = torch.from_numpy(intrinsic1).float()\n intrinsic2 = torch.from_numpy(intrinsic2).float()\n pose = torch.from_numpy(relative[:3, :]).float()\n im1_tensor = self.transform(im1)\n im2_tensor = self.transform(im2)\n\n out = {'im1': im1_tensor,\n 'im2': im2_tensor,\n 'im1_ori': im1_ori,\n 'im2_ori': im2_ori,\n 'pose': pose,\n 'F': F_gt,\n 'intrinsic1': intrinsic1,\n 'intrinsic2': intrinsic2,\n 'coord1': coord1}\n\n return out\n\n def __len__(self):\n return len(self.imf1s)\n" }, { "path": "environment.yml", "content": "name: caps\nchannels:\n - pytorch\n - defaults\ndependencies:\n - _libgcc_mutex=0.1=main\n - blas=1.0=mkl\n - ca-certificates=2020.7.22=0\n - certifi=2020.6.20=py37_0\n - cffi=1.14.1=py37he30daa8_0\n - cudatoolkit=10.0.130=0\n - freetype=2.10.2=h5ab3b9f_0\n - intel-openmp=2020.1=217\n - jpeg=9b=h024ee3a_2\n - lcms2=2.11=h396b838_0\n - ld_impl_linux-64=2.33.1=h53a641e_7\n - libedit=3.1.20191231=h14c3975_1\n - libffi=3.3=he6710b0_2\n - libgcc-ng=9.1.0=hdf63c60_0\n - libpng=1.6.37=hbc83047_0\n - libstdcxx-ng=9.1.0=hdf63c60_0\n - libtiff=4.1.0=h2733197_1\n - lz4-c=1.9.2=he6710b0_1\n - mkl=2020.1=217\n - mkl-service=2.3.0=py37he904b0f_0\n - mkl_fft=1.1.0=py37h23d657b_0\n - mkl_random=1.1.1=py37h0573a6f_0\n - ncurses=6.2=he6710b0_1\n - ninja=1.10.0=py37hfd86e86_0\n - numpy=1.19.1=py37hbc911f0_0\n - numpy-base=1.19.1=py37hfa32c7d_0\n - olefile=0.46=py37_0\n - openssl=1.1.1g=h7b6447c_0\n - pillow=7.2.0=py37hb39fc2d_0\n - pip=20.2.2=py37_0\n - pycparser=2.20=py_2\n - python=3.7.7=hcff3b4d_5\n - pytorch=1.0.1=py3.7_cuda10.0.130_cudnn7.4.2_2\n - readline=8.0=h7b6447c_0\n - setuptools=49.6.0=py37_0\n - six=1.15.0=py_0\n - sqlite=3.33.0=h62c20be_0\n - tk=8.6.10=hbc83047_0\n - torchvision=0.2.2=py_3\n - wheel=0.34.2=py37_0\n - xz=5.2.5=h7b6447c_0\n - zlib=1.2.11=h7b6447c_3\n - zstd=1.4.5=h9ceee32_0\n - pip:\n - chardet==3.0.4\n - configargparse==1.2.3\n - cycler==0.10.0\n - decorator==4.4.2\n - filelock==3.0.12\n - gdown==3.12.2\n - idna==2.10\n - imageio==2.9.0\n - kiwisolver==1.2.0\n - matplotlib==3.3.1\n - networkx==2.5\n - opencv-contrib-python==3.4.2.17\n - opencv-python==3.4.2.17\n - protobuf==3.13.0\n - pyparsing==2.4.7\n - pysocks==1.7.1\n - python-dateutil==2.8.1\n - pywavelets==1.1.1\n - pyyaml==5.3.1\n - requests==2.24.0\n - scikit-image==0.15.0\n - scipy==1.3.1\n - tensorboardx==2.1\n - tqdm==4.48.2\n - urllib3==1.25.10\nprefix: /phoenix/S7/qw246/anaconda3/envs/caps\n\n" }, { "path": "extract_features.py", "content": "import torch\nfrom torch.utils.data import Dataset\nimport os\nimport numpy as np\nimport cv2\nimport skimage.io as io\nimport torchvision.transforms as transforms\nimport config\nfrom tqdm import tqdm\nfrom CAPS.caps_model import CAPSModel\n\n\nclass HPatchDataset(Dataset):\n def __init__(self, imdir):\n self.transform = transforms.Compose([transforms.ToTensor(),\n transforms.Normalize(mean=(0.485, 0.456, 0.406),\n std=(0.229, 0.224, 0.225)),\n ])\n self.imfs = []\n for f in os.listdir(imdir):\n scene_dir = os.path.join(imdir, f)\n self.imfs.extend([os.path.join(scene_dir, '{}.ppm').format(ind) for ind in range(1, 7)])\n\n def __getitem__(self, item):\n imf = self.imfs[item]\n im = io.imread(imf)\n im_tensor = self.transform(im)\n # using sift keypoints\n sift = cv2.xfeatures2d.SIFT_create()\n gray = cv2.cvtColor(im, cv2.COLOR_RGB2GRAY)\n kpts = sift.detect(gray)\n kpts = np.array([[kp.pt[0], kp.pt[1]] for kp in kpts])\n coord = torch.from_numpy(kpts).float()\n out = {'im': im_tensor, 'coord': coord, 'imf': imf}\n return out\n\n def __len__(self):\n return len(self.imfs)\n\n\nif __name__ == '__main__':\n # example code for extracting features for HPatches dataset, SIFT keypoint is used\n args = config.get_args()\n device = 'cuda' if torch.cuda.is_available() else 'cpu'\n\n dataset = HPatchDataset(args.extract_img_dir)\n data_loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=args.workers)\n\n model = CAPSModel(args)\n\n outdir = args.extract_out_dir\n os.makedirs(outdir, exist_ok=True)\n\n with torch.no_grad():\n for data in tqdm(data_loader):\n im = data['im'].to(device)\n img_path = data['imf'][0]\n coord = data['coord'].to(device)\n feat_c, feat_f = model.extract_features(im, coord)\n desc = torch.cat((feat_c, feat_f), -1).squeeze(0).detach().cpu().numpy()\n kpt = coord.cpu().numpy().squeeze(0)\n\n out_path = os.path.join(outdir, '{}-{}'.format(os.path.basename(os.path.dirname(img_path)),\n os.path.basename(img_path),\n ))\n with open(out_path + '.caps', 'wb') as output_file:\n np.savez(\n output_file,\n keypoints=kpt,\n scores=[],\n descriptors=desc\n )\n\n" }, { "path": "jupyter/functions.py", "content": "import numpy as np\nimport matplotlib.pyplot as plt\nimport matplotlib.lines as mlines\nimport torch\nimport sys\nsys.path.append('../')\nfrom dataloader import megadepth\nimport torch.utils.data\nfrom CAPS.caps_model import CAPSModel\nimport cv2\n\n\nclass Visualization(object):\n def __init__(self, args):\n dataset = megadepth.MegaDepth(args)\n self.dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=1)\n self.model = CAPSModel(args)\n self.loader_iter = iter(self.dataloader)\n\n def random_sample(self):\n self.sample = next(self.loader_iter)\n\n def plot_img_pair(self, with_std=False, with_epipline=False):\n self.coords = []\n self.colors = []\n self.with_std = with_std\n self.with_epipline = with_epipline\n im1 = self.sample['im1_ori']\n im2 = self.sample['im2_ori']\n self.h, self.w = im1.shape[1], im1.shape[2]\n im1 = im1.squeeze().cpu().numpy()\n im2 = im2.squeeze().cpu().numpy()\n blank = np.ones((self.h, 5, 3)) * 255\n out = np.concatenate((im1, blank, im2), 1).astype(np.uint8)\n self.fig = plt.figure(figsize=(12, 5))\n self.ax = self.fig.add_subplot(111)\n self.ax.imshow(out)\n self.ax.axis('off')\n plt.tight_layout()\n cid = self.fig.canvas.mpl_connect('button_press_event', self.onclick)\n\n def onclick(self, event):\n color = tuple(np.random.rand(3).tolist())\n coord = [event.xdata, event.ydata]\n self.coord = coord\n self.color = color\n self.coords.append(coord)\n self.colors.append(color)\n self.ax.scatter(event.xdata, event.ydata, c=color)\n self.find_correspondence()\n self.plot_correspondence()\n\n def find_correspondence(self):\n data_in = self.sample\n data_in['coord1'] = torch.from_numpy(np.array(self.coord)).float().cuda().unsqueeze(0).unsqueeze(0)\n data_in['coord2'] = data_in['coord1']\n self.model.set_input(data_in)\n coord2_e, std = self.model.test()\n self.correspondence = coord2_e.squeeze().cpu().numpy()\n self.std = std.squeeze().cpu().numpy()\n\n def plot_correspondence(self):\n point1 = self.coord\n point2 = self.correspondence\n point2[0] += self.w + 5\n self.ax.scatter(point2[0], point2[1], color=self.color)\n if self.with_std:\n circle = plt.Circle((point2[0], point2[1]), radius=100 * self.std, fill=False, color=self.color)\n self.ax.add_patch(circle)\n\n if self.with_epipline:\n line2 = cv2.computeCorrespondEpilines(np.array(point1).reshape(-1, 1, 2), 1,\n self.sample['F'].squeeze().cpu().numpy())\n line2 = np.array(line2).squeeze()\n intersection = np.array([[0, -line2[2]/line2[1]],\n [-line2[2]/line2[0], 0],\n [self.w-1, -(line2[2]+line2[0]*(self.w-1))/line2[1]],\n [-(line2[1]*(self.h-1)+line2[2])/line2[0], self.h-1]])\n valid = (intersection[:, 0] >= 0) & (intersection[:, 0] <= self.w-1) & \\\n (intersection[:, 1] >= 0) & (intersection[:, 1] <= self.h-1)\n if np.sum(valid) == 2:\n intersection = intersection[valid].astype(int)\n x0, y0 = intersection[0]\n x1, y1 = intersection[1]\n l = mlines.Line2D([x0+self.w + 5, x1+self.w + 5], [y0, y1], color=self.color)\n self.ax.add_line(l)\n\n plt.show()\n\n\n\n\n\n\n\n\n\n" }, { "path": "jupyter/visualization.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Interactive demo\\n\",\n \"\\n\",\n \"This interactive demo shows the correspondences our method obtain by **densely** searching in the image space.\\n\",\n \"\\n\",\n \"\\n\",\n \"Run the first code block and wait for it to complete ...\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 25,\n \"metadata\": {\n \"scrolled\": true\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"The autoreload extension is already loaded. To reload it, use:\\n\",\n \" %reload_ext autoreload\\n\"\n ]\n },\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\r\",\n \"# loading data from scene folders: 0%| | 0/37 [00:00');\\n\",\n \" this._root_extra_style(this.root)\\n\",\n \" this.root.attr('style', 'display: inline-block');\\n\",\n \"\\n\",\n \" $(parent_element).append(this.root);\\n\",\n \"\\n\",\n \" this._init_header(this);\\n\",\n \" this._init_canvas(this);\\n\",\n \" this._init_toolbar(this);\\n\",\n \"\\n\",\n \" var fig = this;\\n\",\n \"\\n\",\n \" this.waiting = false;\\n\",\n \"\\n\",\n \" this.ws.onopen = function () {\\n\",\n \" fig.send_message(\\\"supports_binary\\\", {value: fig.supports_binary});\\n\",\n \" fig.send_message(\\\"send_image_mode\\\", {});\\n\",\n \" if (mpl.ratio != 1) {\\n\",\n \" fig.send_message(\\\"set_dpi_ratio\\\", {'dpi_ratio': mpl.ratio});\\n\",\n \" }\\n\",\n \" fig.send_message(\\\"refresh\\\", {});\\n\",\n \" }\\n\",\n \"\\n\",\n \" this.imageObj.onload = function() {\\n\",\n \" if (fig.image_mode == 'full') {\\n\",\n \" // Full images could contain transparency (where diff images\\n\",\n \" // almost always do), so we need to clear the canvas so that\\n\",\n \" // there is no ghosting.\\n\",\n \" fig.context.clearRect(0, 0, fig.canvas.width, fig.canvas.height);\\n\",\n \" }\\n\",\n \" fig.context.drawImage(fig.imageObj, 0, 0);\\n\",\n \" };\\n\",\n \"\\n\",\n \" this.imageObj.onunload = function() {\\n\",\n \" fig.ws.close();\\n\",\n \" }\\n\",\n \"\\n\",\n \" this.ws.onmessage = this._make_on_message_function(this);\\n\",\n \"\\n\",\n \" this.ondownload = ondownload;\\n\",\n \"}\\n\",\n \"\\n\",\n \"mpl.figure.prototype._init_header = function() {\\n\",\n \" var titlebar = $(\\n\",\n \" '
');\\n\",\n \" var titletext = $(\\n\",\n \" '
');\\n\",\n \" titlebar.append(titletext)\\n\",\n \" this.root.append(titlebar);\\n\",\n \" this.header = titletext[0];\\n\",\n \"}\\n\",\n \"\\n\",\n \"\\n\",\n \"\\n\",\n \"mpl.figure.prototype._canvas_extra_style = function(canvas_div) {\\n\",\n \"\\n\",\n \"}\\n\",\n \"\\n\",\n \"\\n\",\n \"mpl.figure.prototype._root_extra_style = function(canvas_div) {\\n\",\n \"\\n\",\n \"}\\n\",\n \"\\n\",\n \"mpl.figure.prototype._init_canvas = function() {\\n\",\n \" var fig = this;\\n\",\n \"\\n\",\n \" var canvas_div = $('
');\\n\",\n \"\\n\",\n \" canvas_div.attr('style', 'position: relative; clear: both; outline: 0');\\n\",\n \"\\n\",\n \" function canvas_keyboard_event(event) {\\n\",\n \" return fig.key_event(event, event['data']);\\n\",\n \" }\\n\",\n \"\\n\",\n \" canvas_div.keydown('key_press', canvas_keyboard_event);\\n\",\n \" canvas_div.keyup('key_release', canvas_keyboard_event);\\n\",\n \" this.canvas_div = canvas_div\\n\",\n \" this._canvas_extra_style(canvas_div)\\n\",\n \" this.root.append(canvas_div);\\n\",\n \"\\n\",\n \" var canvas = $('');\\n\",\n \" canvas.addClass('mpl-canvas');\\n\",\n \" canvas.attr('style', \\\"left: 0; top: 0; z-index: 0; outline: 0\\\")\\n\",\n \"\\n\",\n \" this.canvas = canvas[0];\\n\",\n \" this.context = canvas[0].getContext(\\\"2d\\\");\\n\",\n \"\\n\",\n \" var backingStore = this.context.backingStorePixelRatio ||\\n\",\n \"\\tthis.context.webkitBackingStorePixelRatio ||\\n\",\n \"\\tthis.context.mozBackingStorePixelRatio ||\\n\",\n \"\\tthis.context.msBackingStorePixelRatio ||\\n\",\n \"\\tthis.context.oBackingStorePixelRatio ||\\n\",\n \"\\tthis.context.backingStorePixelRatio || 1;\\n\",\n \"\\n\",\n \" mpl.ratio = (window.devicePixelRatio || 1) / backingStore;\\n\",\n \"\\n\",\n \" var rubberband = $('');\\n\",\n \" rubberband.attr('style', \\\"position: absolute; left: 0; top: 0; z-index: 1;\\\")\\n\",\n \"\\n\",\n \" var pass_mouse_events = true;\\n\",\n \"\\n\",\n \" canvas_div.resizable({\\n\",\n \" start: function(event, ui) {\\n\",\n \" pass_mouse_events = false;\\n\",\n \" },\\n\",\n \" resize: function(event, ui) {\\n\",\n \" fig.request_resize(ui.size.width, ui.size.height);\\n\",\n \" },\\n\",\n \" stop: function(event, ui) {\\n\",\n \" pass_mouse_events = true;\\n\",\n \" fig.request_resize(ui.size.width, ui.size.height);\\n\",\n \" },\\n\",\n \" });\\n\",\n \"\\n\",\n \" function mouse_event_fn(event) {\\n\",\n \" if (pass_mouse_events)\\n\",\n \" return fig.mouse_event(event, event['data']);\\n\",\n \" }\\n\",\n \"\\n\",\n \" rubberband.mousedown('button_press', mouse_event_fn);\\n\",\n \" rubberband.mouseup('button_release', mouse_event_fn);\\n\",\n \" // Throttle sequential mouse events to 1 every 20ms.\\n\",\n \" rubberband.mousemove('motion_notify', mouse_event_fn);\\n\",\n \"\\n\",\n \" rubberband.mouseenter('figure_enter', mouse_event_fn);\\n\",\n \" rubberband.mouseleave('figure_leave', mouse_event_fn);\\n\",\n \"\\n\",\n \" canvas_div.on(\\\"wheel\\\", function (event) {\\n\",\n \" event = event.originalEvent;\\n\",\n \" event['data'] = 'scroll'\\n\",\n \" if (event.deltaY < 0) {\\n\",\n \" event.step = 1;\\n\",\n \" } else {\\n\",\n \" event.step = -1;\\n\",\n \" }\\n\",\n \" mouse_event_fn(event);\\n\",\n \" });\\n\",\n \"\\n\",\n \" canvas_div.append(canvas);\\n\",\n \" canvas_div.append(rubberband);\\n\",\n \"\\n\",\n \" this.rubberband = rubberband;\\n\",\n \" this.rubberband_canvas = rubberband[0];\\n\",\n \" this.rubberband_context = rubberband[0].getContext(\\\"2d\\\");\\n\",\n \" this.rubberband_context.strokeStyle = \\\"#000000\\\";\\n\",\n \"\\n\",\n \" this._resize_canvas = function(width, height) {\\n\",\n \" // Keep the size of the canvas, canvas container, and rubber band\\n\",\n \" // canvas in synch.\\n\",\n \" canvas_div.css('width', width)\\n\",\n \" canvas_div.css('height', height)\\n\",\n \"\\n\",\n \" canvas.attr('width', width * mpl.ratio);\\n\",\n \" canvas.attr('height', height * mpl.ratio);\\n\",\n \" canvas.attr('style', 'width: ' + width + 'px; height: ' + height + 'px;');\\n\",\n \"\\n\",\n \" rubberband.attr('width', width);\\n\",\n \" rubberband.attr('height', height);\\n\",\n \" }\\n\",\n \"\\n\",\n \" // Set the figure to an initial 600x600px, this will subsequently be updated\\n\",\n \" // upon first draw.\\n\",\n \" this._resize_canvas(600, 600);\\n\",\n \"\\n\",\n \" // Disable right mouse context menu.\\n\",\n \" $(this.rubberband_canvas).bind(\\\"contextmenu\\\",function(e){\\n\",\n \" return false;\\n\",\n \" });\\n\",\n \"\\n\",\n \" function set_focus () {\\n\",\n \" canvas.focus();\\n\",\n \" canvas_div.focus();\\n\",\n \" }\\n\",\n \"\\n\",\n \" window.setTimeout(set_focus, 100);\\n\",\n \"}\\n\",\n \"\\n\",\n \"mpl.figure.prototype._init_toolbar = function() {\\n\",\n \" var fig = this;\\n\",\n \"\\n\",\n \" var nav_element = $('
');\\n\",\n \" nav_element.attr('style', 'width: 100%');\\n\",\n \" this.root.append(nav_element);\\n\",\n \"\\n\",\n \" // Define a callback function for later on.\\n\",\n \" function toolbar_event(event) {\\n\",\n \" return fig.toolbar_button_onclick(event['data']);\\n\",\n \" }\\n\",\n \" function toolbar_mouse_event(event) {\\n\",\n \" return fig.toolbar_button_onmouseover(event['data']);\\n\",\n \" }\\n\",\n \"\\n\",\n \" for(var toolbar_ind in mpl.toolbar_items) {\\n\",\n \" var name = mpl.toolbar_items[toolbar_ind][0];\\n\",\n \" var tooltip = mpl.toolbar_items[toolbar_ind][1];\\n\",\n \" var image = mpl.toolbar_items[toolbar_ind][2];\\n\",\n \" var method_name = mpl.toolbar_items[toolbar_ind][3];\\n\",\n \"\\n\",\n \" if (!name) {\\n\",\n \" // put a spacer in here.\\n\",\n \" continue;\\n\",\n \" }\\n\",\n \" var button = $('