added CHM classification

CHMCorr.py

@@ -0,0 +1,546 @@
+# CHM-Corr Classifier
+import argparse
+import json
+import pickle
+import random
+from itertools import product
+import numpy as np
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
+from torch.utils.data import DataLoader
+from torchvision.datasets import ImageFolder
+from tqdm import tqdm
+from common.evaluation import Evaluator
+from model import chmnet
+from model.base.geometry import Geometry
+
+from Utils import (
+    CosineCustomDataset,
+    PairedLayer4Extractor,
+    compute_spatial_similarity,
+    generate_mask,
+    normalize_array,
+    get_transforms,
+    arg_topK,
+)
+
+# Setting the random seed
+random.seed(42)
+
+# Helper Function
+to_np = lambda x: x.data.to("cpu").numpy()
+
+# CHMNet Config
+chm_args = dict(
+    {
+        "alpha": [0.05, 0.1],
+        "img_size": 240,
+        "ktype": "psi",
+        "load": "pas_psi.pt",
+    }
+)
+
+
+class CHMGridTransfer:
+    def __init__(
+        self,
+        query_image,
+        support_set,
+        support_set_labels,
+        train_folder,
+        top_N,
+        top_K,
+        binarization_threshold,
+        chm_source_transform,
+        chm_target_transform,
+        cosine_source_transform,
+        cosine_target_transform,
+        batch_size=64,
+    ):
+        self.N = top_N
+        self.K = top_K
+        self.BS = batch_size
+
+        self.chm_source_transform = chm_source_transform
+        self.chm_target_transform = chm_target_transform
+        self.cosine_source_transform = cosine_source_transform
+        self.cosine_target_transform = cosine_target_transform
+
+        self.source_embeddings = None
+        self.target_embeddings = None
+        self.correspondence_map = None
+        self.similarity_maps = None
+        self.reverse_similarity_maps = None
+        self.transferred_points = None
+
+        self.binarization_threshold = binarization_threshold
+
+        # UPDATE THIS
+        self.q = query_image
+        self.support_set = support_set
+        self.labels_ss = support_set_labels
+
+    def build(self):
+        # C.M.H
+        test_ds = CosineCustomDataset(
+            query_image=self.q,
+            supporting_set=self.support_set,
+            source_transform=self.chm_source_transform,
+            target_transform=self.chm_target_transform,
+        )
+        test_dl = DataLoader(test_ds, batch_size=self.BS, shuffle=False)
+        self.find_correspondences(test_dl)
+
+        # LAYER 4s
+        test_ds = CosineCustomDataset(
+            query_image=self.q,
+            supporting_set=self.support_set,
+            source_transform=self.cosine_source_transform,
+            target_transform=self.cosine_target_transform,
+        )
+        test_dl = DataLoader(test_ds, batch_size=self.BS, shuffle=False)
+        self.compute_embeddings(test_dl)
+        self.compute_similarity_map()
+
+    def find_correspondences(self, test_dl):
+        model = chmnet.CHMNet(chm_args["ktype"])
+        model.load_state_dict(
+            torch.load(chm_args["load"], map_location=torch.device("cpu"))
+        )
+        Evaluator.initialize(chm_args["alpha"])
+        Geometry.initialize(img_size=chm_args["img_size"])
+
+        grid_results = []
+        transferred_points = []
+
+        # FIXED GRID HARD CODED
+        fixed_src_grid_points = list(
+            product(
+                np.linspace(1 + 17, 240 - 17 - 1, 7),
+                np.linspace(1 + 17, 240 - 17 - 1, 7),
+            )
+        )
+        fixed_src_grid_points = np.asarray(fixed_src_grid_points, dtype=np.float64).T
+
+        with torch.no_grad():
+            model.eval()
+            for idx, batch in enumerate(tqdm(test_dl)):
+
+                keypoints = (
+                    torch.tensor(fixed_src_grid_points)
+                    .unsqueeze(0)
+                    .repeat(batch["src_img"].shape[0], 1, 1)
+                )
+                n_pts = torch.tensor(
+                    np.asarray(batch["src_img"].shape[0] * [49]), dtype=torch.long
+                )
+
+                corr_matrix = model(batch["src_img"], batch["trg_img"])
+                prd_kps = Geometry.transfer_kps(
+                    corr_matrix, keypoints, n_pts, normalized=False
+                )
+                transferred_points.append(prd_kps.cpu().numpy())
+                for tgt_points in prd_kps:
+                    tgt_grid = []
+                    for x, y in zip(tgt_points[0], tgt_points[1]):
+                        tgt_grid.append(
+                            [int(((x + 1) / 2.0) * 7), int(((y + 1) / 2.0) * 7)]
+                        )
+                    grid_results.append(tgt_grid)
+
+        self.correspondence_map = grid_results
+        self.transferred_points = np.vstack(transferred_points)
+
+    def compute_embeddings(self, test_dl):
+        paired_extractor = PairedLayer4Extractor()
+
+        source_embeddings = []
+        target_embeddings = []
+
+        with torch.no_grad():
+            for idx, batch in enumerate(test_dl):
+                s_e, t_e = paired_extractor((batch["src_img"], batch["trg_img"]))
+
+                source_embeddings.append(s_e)
+                target_embeddings.append(t_e)
+
+        # EMBEDDINGS
+        self.source_embeddings = torch.cat(source_embeddings, axis=0)
+        self.target_embeddings = torch.cat(target_embeddings, axis=0)
+
+    def compute_similarity_map(self):
+        CosSim = nn.CosineSimilarity(dim=0, eps=1e-6)
+
+        similarity_maps = []
+        rsimilarity_maps = []
+
+        grid = []
+        for i in range(7):
+            for j in range(7):
+                grid.append([i, j])
+
+        # Compute for all image pairs
+        for i in range(len(self.correspondence_map)):
+            cosine_map = np.zeros((7, 7))
+            reverse_cosine_map = np.zeros((7, 7))
+
+            # calculate cosine based on the chm corr. map
+            for S, T in zip(grid, self.correspondence_map[i]):
+                v1 = self.source_embeddings[i][:, S[0], S[1]]
+                v2 = self.target_embeddings[i][:, T[0], T[1]]
+                covalue = CosSim(v1, v2)
+                cosine_map[S[0], S[1]] = covalue
+                reverse_cosine_map[T[0], T[1]] = covalue
+
+            similarity_maps.append(cosine_map)
+            rsimilarity_maps.append(reverse_cosine_map)
+
+        self.similarity_maps = similarity_maps
+        self.reverse_similarity_maps = rsimilarity_maps
+
+    def compute_score_using_cc(self):
+        # CC MAPS
+        SIMS_source, SIMS_target = [], []
+        for i in range(len(self.source_embeddings)):
+            simA, simB = compute_spatial_similarity(
+                to_np(self.source_embeddings[i]), to_np(self.target_embeddings[i])
+            )
+
+            SIMS_source.append(simA)
+            SIMS_target.append(simB)
+
+        SIMS_source = np.stack(SIMS_source, axis=0)
+        # SIMS_target = np.stack(SIMS_target, axis=0)
+
+        top_cos_values = []
+
+        for i in range(len(self.similarity_maps)):
+            cosine_value = np.multiply(
+                self.similarity_maps[i],
+                generate_mask(
+                    normalize_array(SIMS_source[i]), t=self.binarization_threshold
+                ),
+            )
+            top_5_indicies = np.argsort(cosine_value.T.reshape(-1))[::-1][:5]
+            mean_of_top_5 = np.mean(
+                [cosine_value.T.reshape(-1)[x] for x in top_5_indicies]
+            )
+            top_cos_values.append(np.mean(mean_of_top_5))
+
+        return top_cos_values
+
+    def compute_score_using_custom_points(self, selected_keypoint_masks):
+        top_cos_values = []
+
+        for i in range(len(self.similarity_maps)):
+            cosine_value = np.multiply(self.similarity_maps[i], selected_keypoint_masks)
+            top_indicies = np.argsort(cosine_value.T.reshape(-1))[::-1]
+            mean_of_tops = np.mean(
+                [cosine_value.T.reshape(-1)[x] for x in top_indicies]
+            )
+            top_cos_values.append(np.mean(mean_of_tops))
+
+        return top_cos_values
+
+    def export(self):
+        storage = {
+            "N": self.N,
+            "K": self.K,
+            "source_embeddings": self.source_embeddings,
+            "target_embeddings": self.target_embeddings,
+            "correspondence_map": self.correspondence_map,
+            "similarity_maps": self.similarity_maps,
+            "T": self.binarization_threshold,
+            "query": self.q,
+            "support_set": self.support_set,
+            "labels_for_support_set": self.labels_ss,
+            "rsimilarity_maps": self.reverse_similarity_maps,
+            "transferred_points": self.transferred_points,
+        }
+
+        return ModifiableCHMResults(storage)
+
+
+class ModifiableCHMResults:
+    def __init__(self, storage):
+        self.N = storage["N"]
+        self.K = storage["K"]
+        self.source_embeddings = storage["source_embeddings"]
+        self.target_embeddings = storage["target_embeddings"]
+        self.correspondence_map = storage["correspondence_map"]
+        self.similarity_maps = storage["similarity_maps"]
+        self.T = storage["T"]
+        self.q = storage["query"]
+        self.support_set = storage["support_set"]
+        self.labels_ss = storage["labels_for_support_set"]
+        self.rsimilarity_maps = storage["rsimilarity_maps"]
+        self.transferred_points = storage["transferred_points"]
+        self.similarity_maps_masked = None
+        self.SIMS_source = None
+        self.SIMS_target = None
+        self.masked_sim_values = []
+        self.top_cos_values = []
+
+    def compute_score_using_cc(self):
+        # CC MAPS
+        SIMS_source, SIMS_target = [], []
+        for i in range(len(self.source_embeddings)):
+            simA, simB = compute_spatial_similarity(
+                to_np(self.source_embeddings[i]), to_np(self.target_embeddings[i])
+            )
+
+            SIMS_source.append(simA)
+            SIMS_target.append(simB)
+
+        SIMS_source = np.stack(SIMS_source, axis=0)
+        SIMS_target = np.stack(SIMS_target, axis=0)
+
+        self.SIMS_source = SIMS_source
+        self.SIMS_target = SIMS_target
+
+        top_cos_values = []
+
+        for i in range(len(self.similarity_maps)):
+            masked_sim_values = np.multiply(
+                self.similarity_maps[i],
+                generate_mask(normalize_array(SIMS_source[i]), t=self.T),
+            )
+            self.masked_sim_values.append(masked_sim_values)
+            top_5_indicies = np.argsort(masked_sim_values.T.reshape(-1))[::-1][:5]
+            mean_of_top_5 = np.mean(
+                [masked_sim_values.T.reshape(-1)[x] for x in top_5_indicies]
+            )
+            top_cos_values.append(np.mean(mean_of_top_5))
+
+        self.top_cos_values = top_cos_values
+
+        return top_cos_values
+
+    def compute_score_using_custom_points(self, selected_keypoint_masks):
+        top_cos_values = []
+        similarity_maps_masked = []
+
+        for i in range(len(self.similarity_maps)):
+            cosine_value = np.multiply(self.similarity_maps[i], selected_keypoint_masks)
+            similarity_maps_masked.append(cosine_value)
+            top_indicies = np.argsort(cosine_value.T.reshape(-1))[::-1]
+            mean_of_tops = np.mean(
+                [cosine_value.T.reshape(-1)[x] for x in top_indicies]
+            )
+            top_cos_values.append(np.mean(mean_of_tops))
+
+        self.similarity_maps_masked = similarity_maps_masked
+        return top_cos_values
+
+    def predict_using_cc(self):
+        top_cos_values = self.compute_score_using_cc()
+        # Predict
+        prediction = np.argmax(
+            np.bincount(
+                [self.labels_ss[x] for x in np.argsort(top_cos_values)[::-1][: self.K]]
+            )
+        )
+        prediction_weight = np.max(
+            np.bincount(
+                [self.labels_ss[x] for x in np.argsort(top_cos_values)[::-1][: self.K]]
+            )
+        )
+
+        reranked_nns_idx = [x for x in np.argsort(top_cos_values)[::-1]]
+        reranked_nns_files = [self.support_set[x] for x in reranked_nns_idx]
+
+        topK_idx = [
+            x
+            for x in np.argsort(top_cos_values)[::-1]
+            if self.labels_ss[x] == prediction
+        ]
+        topK_files = [self.support_set[x] for x in topK_idx]
+        topK_cmaps = [self.correspondence_map[x] for x in topK_idx]
+        topK_similarity_maps = [self.similarity_maps[x] for x in topK_idx]
+        topK_rsimilarity_maps = [self.rsimilarity_maps[x] for x in topK_idx]
+        topK_transfered_points = [self.transferred_points[x] for x in topK_idx]
+        predicted_folder_name = topK_files[0].split("/")[-2]
+
+        return (
+            topK_idx,
+            prediction,
+            predicted_folder_name,
+            prediction_weight,
+            topK_files[: self.K],
+            reranked_nns_files[: self.K],
+            topK_cmaps[: self.K],
+            topK_similarity_maps[: self.K],
+            topK_rsimilarity_maps[: self.K],
+            topK_transfered_points[: self.K],
+        )
+
+    def predict_custom_pairs(self, selected_keypoint_masks):
+        top_cos_values = self.compute_score_using_custom_points(selected_keypoint_masks)
+
+        # Predict
+        prediction = np.argmax(
+            np.bincount(
+                [self.labels_ss[x] for x in np.argsort(top_cos_values)[::-1][: self.K]]
+            )
+        )
+        prediction_weight = np.max(
+            np.bincount(
+                [self.labels_ss[x] for x in np.argsort(top_cos_values)[::-1][: self.K]]
+            )
+        )
+
+        reranked_nns_idx = [x for x in np.argsort(top_cos_values)[::-1]]
+        reranked_nns_files = [self.support_set[x] for x in reranked_nns_idx]
+
+        topK_idx = [
+            x
+            for x in np.argsort(top_cos_values)[::-1]
+            if self.labels_ss[x] == prediction
+        ]
+        topK_files = [self.support_set[x] for x in topK_idx]
+        topK_cmaps = [self.correspondence_map[x] for x in topK_idx]
+        topK_similarity_maps = [self.similarity_maps[x] for x in topK_idx]
+        topK_rsimilarity_maps = [self.rsimilarity_maps[x] for x in topK_idx]
+        topK_transferred_points = [self.transferred_points[x] for x in topK_idx]
+        # topK_scores = [top_cos_values[x] for x in topK_idx]
+        topK_masked_sims = [self.similarity_maps_masked[x] for x in topK_idx]
+        predicted_folder_name = topK_files[0].split("/")[-2]
+
+        non_zero_mask = np.count_nonzero(selected_keypoint_masks)
+
+        return (
+            topK_idx,
+            prediction,
+            predicted_folder_name,
+            prediction_weight,
+            topK_files[: self.K],
+            reranked_nns_files[: self.K],
+            topK_cmaps[: self.K],
+            topK_similarity_maps[: self.K],
+            topK_rsimilarity_maps[: self.K],
+            topK_transferred_points[: self.K],
+            topK_masked_sims[: self.K],
+            non_zero_mask,
+        )
+
+
+def export_visualizations_results(
+    reranker_output,
+    knn_predicted_label,
+    knn_confidence,
+    topK_knns,
+    K=20,
+    N=50,
+    T=0.55,
+):
+    """
+    Export all details for visualization and analysis
+    """
+
+    non_zero_mask = 5  # default value
+    (
+        topK_idx,
+        p,
+        pfn,
+        pr,
+        rfiles,
+        reranked_nns,
+        cmaps,
+        sims,
+        rsims,
+        trns_kpts,
+    ) = reranker_output.predict_using_cc()
+
+    MASKED_COSINE_VALUES = [
+        np.multiply(
+            sims[X],
+            generate_mask(
+                normalize_array(reranker_output.SIMS_source[topK_idx[X]]), t=T
+            ),
+        )
+        for X in range(len(sims))
+    ]
+
+    list_of_source_points = []
+    list_of_target_points = []
+
+    for CK in range(len(sims)):
+        target_keypoints = []
+        topk_index = arg_topK(MASKED_COSINE_VALUES[CK], topK=non_zero_mask)
+
+        for i in range(non_zero_mask):  # Number of Connections
+            # Psource = point_list[topk_index[i]]
+            x, y = trns_kpts[CK].T[topk_index[i]]
+            Ptarget = int(((x + 1) / 2.0) * 240), int(((y + 1) / 2.0) * 240)
+            target_keypoints.append(Ptarget)
+
+        # Uniform Grid of points
+        a = np.linspace(1 + 17, 240 - 17 - 1, 7)
+        b = np.linspace(1 + 17, 240 - 17 - 1, 7)
+        point_list = list(product(a, b))
+
+        list_of_source_points.append(np.asarray([point_list[x] for x in topk_index]))
+        list_of_target_points.append(np.asarray(target_keypoints))
+
+    # EXPORT OUTPUT
+    detailed_output = {
+        "q": reranker_output.q,
+        "K": K,
+        "N": N,
+        "knn-prediction": knn_predicted_label,
+        "knn-prediction-confidence": knn_confidence,
+        "knn-nearest-neighbors": topK_knns,
+        "chm-prediction": pfn,
+        "chm-prediction-confidence": pr,
+        "chm-nearest-neighbors": rfiles,
+        "correspondance_map": cmaps,
+        "masked_cos_values": MASKED_COSINE_VALUES,
+        "src-keypoints": list_of_source_points,
+        "tgt-keypoints": list_of_target_points,
+        "non_zero_mask": non_zero_mask,
+        "transferred_kpoints": trns_kpts,
+    }
+
+    return detailed_output
+
+
+def chm_classify_and_visualize(
+    query_image, kNN_results, support, TRAIN_SET, N=50, K=20, T=0.55, BS=64
+):
+    global chm_args
+    chm_src_t, chm_tgt_t, cos_src_t, cos_tgt_t = get_transforms("single", chm_args)
+    knn_predicted_label, knn_confidence, topK_knns = kNN_results
+
+    reranker = CHMGridTransfer(
+        query_image=query_image,
+        support_set=support[0],
+        support_set_labels=support[1],
+        train_folder=TRAIN_SET,
+        top_N=N,
+        top_K=K,
+        binarization_threshold=T,
+        chm_source_transform=chm_src_t,
+        chm_target_transform=chm_tgt_t,
+        cosine_source_transform=cos_src_t,
+        cosine_target_transform=cos_tgt_t,
+        batch_size=BS,
+    )
+
+    # Building the reranker
+    reranker.build()
+    # Make a ModifiableCHMResults
+    exported_reranker = reranker.export()
+    # Export A details for visualizations
+
+    output = export_visualizations_results(
+        exported_reranker,
+        knn_predicted_label,
+        knn_confidence,
+        topK_knns,
+        K,
+        N,
+        T,
+    )
+
+    return output

ExtractEmbedding.py

@@ -36,7 +36,7 @@ class Wrapper(torch.nn.Module):
         return "Wrappper"
 
 
-def QueryToEmbedding(query_pil):
+def QueryToEmbedding(query_path):
     dataset_transform = transforms.Compose(
         [
             transforms.Resize(256),
@@ -50,7 +50,7 @@ def QueryToEmbedding(query_pil):
     model.eval()
     myw = Wrapper(model)
 
-    # query_pil = Image.open(query_path)
+    query_pil = Image.open(query_path)
     query_pt = dataset_transform(query_pil)
 
     with torch.no_grad():

FeatureExtractors.py

@@ -0,0 +1,391 @@
+# Original Author: Jonathan Donnellya (jonathan.donnelly@maine.edu)
+# Modified by Mohammad Reza Taesiri (mtaesiri@gmail.com)
+
+import os
+import torch
+import torch.nn as nn
+from collections import OrderedDict
+
+model_dir = os.path.dirname(os.path.realpath(__file__))
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(
+        in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False
+    )
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    # class attribute
+    expansion = 1
+    num_layers = 2
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        # only conv with possibly not 1 stride
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        # if stride is not 1 then self.downsample cannot be None
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        # the residual connection
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+    def block_conv_info(self):
+        block_kernel_sizes = [3, 3]
+        block_strides = [self.stride, 1]
+        block_paddings = [1, 1]
+
+        return block_kernel_sizes, block_strides, block_paddings
+
+
+class Bottleneck(nn.Module):
+    # class attribute
+    expansion = 4
+    num_layers = 3
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = conv1x1(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        # only conv with possibly not 1 stride
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = conv1x1(planes, planes * self.expansion)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+
+        # if stride is not 1 then self.downsample cannot be None
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+    def block_conv_info(self):
+        block_kernel_sizes = [1, 3, 1]
+        block_strides = [1, self.stride, 1]
+        block_paddings = [0, 1, 0]
+
+        return block_kernel_sizes, block_strides, block_paddings
+
+
+class ResNet_features(nn.Module):
+    """
+    the convolutional layers of ResNet
+    the average pooling and final fully convolutional layer is removed
+    """
+
+    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False):
+        super(ResNet_features, self).__init__()
+
+        self.inplanes = 64
+
+        # the first convolutional layer before the structured sequence of blocks
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        # comes from the first conv and the following max pool
+        self.kernel_sizes = [7, 3]
+        self.strides = [2, 2]
+        self.paddings = [3, 1]
+
+        # the following layers, each layer is a sequence of blocks
+        self.block = block
+        self.layers = layers
+        self.layer1 = self._make_layer(
+            block=block, planes=64, num_blocks=self.layers[0]
+        )
+        self.layer2 = self._make_layer(
+            block=block, planes=128, num_blocks=self.layers[1], stride=2
+        )
+        self.layer3 = self._make_layer(
+            block=block, planes=256, num_blocks=self.layers[2], stride=2
+        )
+        self.layer4 = self._make_layer(
+            block=block, planes=512, num_blocks=self.layers[3], stride=2
+        )
+
+        # initialize the parameters
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, num_blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        # only the first block has downsample that is possibly not None
+        layers.append(block(self.inplanes, planes, stride, downsample))
+
+        self.inplanes = planes * block.expansion
+        for _ in range(1, num_blocks):
+            layers.append(block(self.inplanes, planes))
+
+        # keep track of every block's conv size, stride size, and padding size
+        for each_block in layers:
+            (
+                block_kernel_sizes,
+                block_strides,
+                block_paddings,
+            ) = each_block.block_conv_info()
+            self.kernel_sizes.extend(block_kernel_sizes)
+            self.strides.extend(block_strides)
+            self.paddings.extend(block_paddings)
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        return x
+
+    def conv_info(self):
+        return self.kernel_sizes, self.strides, self.paddings
+
+    def num_layers(self):
+        """
+        the number of conv layers in the network, not counting the number
+        of bypass layers
+        """
+
+        return (
+            self.block.num_layers * self.layers[0]
+            + self.block.num_layers * self.layers[1]
+            + self.block.num_layers * self.layers[2]
+            + self.block.num_layers * self.layers[3]
+            + 1
+        )
+
+    def __repr__(self):
+        template = "resnet{}_features"
+        return template.format(self.num_layers() + 1)
+
+
+def resnet50_features(pretrained=True, inat=True, **kwargs):
+    """Constructs a ResNet-50 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet or iNaturalist
+        pretrained (bool): If True, returns a model pre-trained on iNaturalst; else, ImageNet
+    """
+    model = ResNet_features(Bottleneck, [3, 4, 6, 4], **kwargs)
+    if pretrained:
+        if inat:
+            # print('Loading iNat model')
+            model_dict = torch.load(
+                model_dir
+                + "/../../weights/"
+                + "BBN.iNaturalist2017.res50.90epoch.best_model.pth.pt"
+            )
+        else:
+            raise
+
+        if inat:
+            model_dict.pop("module.classifier.weight")
+            model_dict.pop("module.classifier.bias")
+            for key in list(model_dict.keys()):
+                model_dict[
+                    key.replace("module.backbone.", "")
+                    .replace("cb_block", "layer4.2")
+                    .replace("rb_block", "layer4.3")
+                ] = model_dict.pop(key)
+
+        else:
+            raise
+
+        model.load_state_dict(model_dict, strict=False)
+
+    return model
+
+
+class ResNet_classifier(nn.Module):
+    """
+    A classifier for Deformable ProtoPNet
+    """
+
+    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False):
+        super(ResNet_classifier, self).__init__()
+
+        self.inplanes = 64
+
+        # the first convolutional layer before the structured sequence of blocks
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        # comes from the first conv and the following max pool
+        self.kernel_sizes = [7, 3]
+        self.strides = [2, 2]
+        self.paddings = [3, 1]
+
+        # the following layers, each layer is a sequence of blocks
+        self.block = block
+        self.layers = layers
+        self.layer1 = self._make_layer(
+            block=block, planes=64, num_blocks=self.layers[0]
+        )
+        self.layer2 = self._make_layer(
+            block=block, planes=128, num_blocks=self.layers[1], stride=2
+        )
+        self.layer3 = self._make_layer(
+            block=block, planes=256, num_blocks=self.layers[2], stride=2
+        )
+        self.layer4 = self._make_layer(
+            block=block, planes=512, num_blocks=self.layers[3], stride=2
+        )
+
+        self.classifier = nn.Linear(2048 * 7 * 7, 200)
+
+        # initialize the parameters
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, num_blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        # only the first block has downsample that is possibly not None
+        layers.append(block(self.inplanes, planes, stride, downsample))
+
+        self.inplanes = planes * block.expansion
+        for _ in range(1, num_blocks):
+            layers.append(block(self.inplanes, planes))
+
+        # keep track of every block's conv size, stride size, and padding size
+        for each_block in layers:
+            (
+                block_kernel_sizes,
+                block_strides,
+                block_paddings,
+            ) = each_block.block_conv_info()
+            self.kernel_sizes.extend(block_kernel_sizes)
+            self.strides.extend(block_strides)
+            self.paddings.extend(block_paddings)
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.classifier(torch.flatten(x, start_dim=1))
+        return x
+
+    def conv_info(self):
+        return self.kernel_sizes, self.strides, self.paddings
+
+    def num_layers(self):
+        """
+        the number of conv layers in the network, not counting the number
+        of bypass layers
+        """
+
+        return (
+            self.block.num_layers * self.layers[0]
+            + self.block.num_layers * self.layers[1]
+            + self.block.num_layers * self.layers[2]
+            + self.block.num_layers * self.layers[3]
+            + 1
+        )
+
+    def __repr__(self):
+        template = "resnet{}_features"
+        return template.format(self.num_layers() + 1)

Utils.py

@@ -0,0 +1,323 @@
+import numpy as np
+import torch
+import torchvision.models as models
+from numpy import matlib as mb
+from PIL import Image
+from torch.utils.data import Dataset
+from torchvision.datasets import ImageFolder
+import torchvision.transforms as transforms
+from FeatureExtractors import resnet50_features
+
+to_np = lambda x: x.data.to("cpu").numpy()
+
+
+def compute_spatial_similarity(conv1, conv2):
+    """
+    Takes in the last convolutional layer from two images, computes the pooled output
+    feature, and then generates the spatial similarity map for both images.
+    """
+    conv1 = conv1.reshape(-1, 7 * 7).T
+    conv2 = conv2.reshape(-1, 7 * 7).T
+
+    pool1 = np.mean(conv1, axis=0)
+    pool2 = np.mean(conv2, axis=0)
+    out_sz = (int(np.sqrt(conv1.shape[0])), int(np.sqrt(conv1.shape[0])))
+    conv1_normed = conv1 / np.linalg.norm(pool1) / conv1.shape[0]
+    conv2_normed = conv2 / np.linalg.norm(pool2) / conv2.shape[0]
+    im_similarity = np.zeros((conv1_normed.shape[0], conv1_normed.shape[0]))
+
+    for zz in range(conv1_normed.shape[0]):
+        repPx = mb.repmat(conv1_normed[zz, :], conv1_normed.shape[0], 1)
+        im_similarity[zz, :] = np.multiply(repPx, conv2_normed).sum(axis=1)
+    similarity1 = np.reshape(np.sum(im_similarity, axis=1), out_sz)
+    similarity2 = np.reshape(np.sum(im_similarity, axis=0), out_sz)
+    return similarity1, similarity2
+
+
+def normalize_array(x):
+    x = np.asarray(x).copy()
+    x -= np.min(x)
+    x /= np.max(x)
+    return x
+
+
+def apply_threshold(x, t):
+    x = np.asarray(x).copy()
+    x[x < t] = 0
+    return x
+
+
+def generate_mask(x, t):
+    v = np.zeros_like(x)
+    v[x >= t] = 1
+    return v
+
+
+def get_transforms(args_transform, chm_args):
+    # TRANSFORMS
+    cosine_transform_target = transforms.Compose(
+        [
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
+
+    chm_transform_target = transforms.Compose(
+        [
+            transforms.Resize(chm_args["img_size"]),
+            transforms.CenterCrop((chm_args["img_size"], chm_args["img_size"])),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+
+    if args_transform == "multi":
+        cosine_transform_source = transforms.Compose(
+            [
+                transforms.Resize((224, 224)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+            ]
+        )
+
+        chm_transform_source = transforms.Compose(
+            [
+                transforms.Resize((chm_args["img_size"], chm_args["img_size"])),
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+                ),
+            ]
+        )
+
+    elif args_transform == "single":
+        cosine_transform_source = transforms.Compose(
+            [
+                transforms.Resize(chm_args["img_size"]),
+                transforms.CenterCrop((chm_args["img_size"], chm_args["img_size"])),
+                transforms.Resize((224, 224)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+            ]
+        )
+
+        chm_transform_source = transforms.Compose(
+            [
+                transforms.Resize(chm_args["img_size"]),
+                transforms.CenterCrop((chm_args["img_size"], chm_args["img_size"])),
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+                ),
+            ]
+        )
+
+    return (
+        chm_transform_source,
+        chm_transform_target,
+        cosine_transform_source,
+        cosine_transform_target,
+    )
+
+
+def clamp(x, min_value, max_value):
+    return max(min_value, min(x, max_value))
+
+
+def keep_top5(input_array, K=5):
+    top_5 = np.sort(input_array.reshape(-1))[::-1][K - 1]
+    masked = np.zeros_like(input_array)
+    masked[input_array >= top_5] = 1
+    return masked
+
+
+def arg_topK(input_array, topK=5):
+    return np.argsort(input_array.T.reshape(-1))[::-1][:topK]
+
+
+class KNNSupportSet:
+    def __init__(self, train_folder, val_folder, knn_scores, custom_val_labels=None):
+        self.train_data = ImageFolder(root=train_folder)
+        self.val_data = ImageFolder(root=val_folder)
+        self.knn_scores = knn_scores
+
+        if custom_val_labels is None:
+            self.val_labels = np.asarray([x[1] for x in self.val_data.imgs])
+        else:
+            self.val_labels = custom_val_labels
+
+        self.train_labels = np.asarray([x[1] for x in self.train_data.imgs])
+
+    def get_knn_predictions(self, k=20):
+        knn_predictions = [
+            np.argmax(np.bincount(self.train_labels[self.knn_scores[I][::-1][:k]]))
+            for I in range(len(self.knn_scores))
+        ]
+        knn_accuracy = (
+            100
+            * np.sum((np.asarray(knn_predictions) == self.val_labels))
+            / len(self.val_labels)
+        )
+        return knn_predictions, knn_accuracy
+
+    def get_support_set(self, selected_index, top_N=20):
+        support_set = self.knn_scores[selected_index][-top_N:][::-1]
+        return [self.train_data.imgs[x][0] for x in support_set]
+
+    def get_support_set_labels(self, selected_index, top_N=20):
+        support_set = self.knn_scores[selected_index][-top_N:][::-1]
+        return [self.train_data.imgs[x][1] for x in support_set]
+
+    def get_image_and_label_by_id(self, q_id):
+        q = self.val_data.imgs[q_id][0]
+        ql = self.val_data.imgs[q_id][1]
+        return (q, ql)
+
+    def get_folder_name(self, q_id):
+        q = self.val_data.imgs[q_id][0]
+        return q.split("/")[-2]
+
+    def get_top5_knn(self, query_id, k=20):
+        knn_pred, knn_acc = self.get_knn_predictions(k=k)
+        top_5s_index = np.where(
+            np.equal(
+                self.train_labels[self.knn_scores[query_id][::-1]], knn_pred[query_id]
+            )
+        )[0][:5]
+        top_5s = self.knn_scores[query_id][::-1][top_5s_index]
+        top_5s_files = [self.train_data.imgs[x][0] for x in top_5s]
+        return top_5s_files
+
+    def get_topK_knn(self, query_id, k=20):
+        knn_pred, knn_acc = self.get_knn_predictions(k=k)
+        top_ks_index = np.where(
+            np.equal(
+                self.train_labels[self.knn_scores[query_id][::-1]], knn_pred[query_id]
+            )
+        )[0][:k]
+        top_ks = self.knn_scores[query_id][::-1][top_ks_index]
+        top_ks_files = [self.train_data.imgs[x][0] for x in top_ks]
+        return top_ks_files
+
+    def get_foldername_for_label(self, label):
+        for i in range(len(self.train_data)):
+            if self.train_data.imgs[i][1] == label:
+                return self.train_data.imgs[i][0].split("/")[-2]
+
+    def get_knn_confidence(self, query_id, k=20):
+        return np.max(
+            np.bincount(self.train_labels[self.knn_scores[query_id][::-1][:k]])
+        )
+
+
+class CosineCustomDataset(Dataset):
+    r"""Parent class of PFPascal, PFWillow, and SPair"""
+
+    def __init__(self, query_image, supporting_set, source_transform, target_transform):
+        r"""XAICustomDataset constructor"""
+        super(CosineCustomDataset, self).__init__()
+
+        self.supporting_set = supporting_set
+        self.query_image = [query_image] * len(supporting_set)
+
+        self.source_transform = source_transform
+        self.target_transform = target_transform
+
+    def __len__(self):
+        r"""Returns the number of pairs"""
+        return len(self.supporting_set)
+
+    def __getitem__(self, idx):
+        r"""Constructs and return a batch"""
+
+        # Image name
+        batch = dict()
+        batch["src_imname"] = self.query_image[idx]
+        batch["trg_imname"] = self.supporting_set[idx]
+
+        # Image as numpy (original width, original height)
+        src_pil = self.get_image(self.query_image, idx)
+        trg_pil = self.get_image(self.supporting_set, idx)
+
+        batch["src_imsize"] = src_pil.size
+        batch["trg_imsize"] = trg_pil.size
+
+        # Image as tensor
+        batch["src_img"] = self.source_transform(src_pil)
+        batch["trg_img"] = self.target_transform(trg_pil)
+
+        # Total number of pairs in training split
+        batch["datalen"] = len(self.query_image)
+        return batch
+
+    def get_image(self, image_pathes, idx):
+        r"""Reads PIL image from path"""
+        path = image_pathes[idx]
+        return Image.open(path).convert("RGB")
+
+
+class PairedLayer4Extractor(torch.nn.Module):
+    """
+    Extracting layer-4 embedding for source and target images using ResNet-50 features
+    """
+
+    def __init__(self):
+        super(PairedLayer4Extractor, self).__init__()
+
+        self.modelA = models.resnet50(pretrained=True)
+        self.modelA.eval()
+
+        self.modelB = models.resnet50(pretrained=True)
+        self.modelB.eval()
+
+        self.a_embeddings = None
+        self.b_embeddings = None
+
+        def a_hook(module, input, output):
+            self.a_embeddings = output
+
+        def b_hook(module, input, output):
+            self.b_embeddings = output
+
+        self.modelA._modules.get("layer4").register_forward_hook(a_hook)
+        self.modelB._modules.get("layer4").register_forward_hook(b_hook)
+
+    def forward(self, inputs):
+        inputA, inputB = inputs
+        self.modelA(inputA)
+        self.modelB(inputB)
+
+        return self.a_embeddings, self.b_embeddings
+
+    def __repr__(self):
+        return "PairedLayer4Extractor"
+
+
+class iNaturalistPairedLayer4Extractor(torch.nn.Module):
+    """
+    Extracting layer-4 embedding for source and target images using iNaturalist ResNet-50 features
+    """
+
+    def __init__(self):
+        super(iNaturalistPairedLayer4Extractor, self).__init__()
+
+        self.modelA = resnet50_features(inat=True, pretrained=True)
+        self.modelA.eval()
+
+        self.modelB = resnet50_features(inat=True, pretrained=True)
+        self.modelB.eval()
+
+        self.source_embedding = None
+        self.target_embedding = None
+
+    def forward(self, inputs):
+        source_image, target_image = inputs
+        self.source_embedding = self.modelA(source_image)
+        self.target_embedding = self.modelB(target_image)
+
+        return self.source_embedding, self.target_embedding
+
+    def __repr__(self):
+        return "iNatPairedLayer4Extractor"

app.py

@@ -10,12 +10,21 @@ from torchvision.datasets import ImageFolder
 
 from SimSearch import FaissCosineNeighbors, SearchableTrainingSet
 from ExtractEmbedding import QueryToEmbedding
+from CHMCorr import chm_classify_and_visualize
+from visualization import plot_from_reranker_output
 
 csv.field_size_limit(sys.maxsize)
 
 concat = lambda x: np.concatenate(x, axis=0)
 
-gdown.download(id="116CiA_cXciGSl72tbAUDoN-f1B9Frp89")
+# Embeddings
+gdown.cached_download(
+    url="https://drive.google.com/uc?id=116CiA_cXciGSl72tbAUDoN-f1B9Frp89",
+    path="./embeddings.pkl",
+    quiet=False,
+    md5="002b2a7f5c80d910b9cc740c2265f058",
+)
+
 gdown.download(id="1SDtq6ap7LPPpYfLbAxaMGGmj0EAV_m_e")
 
 # CUB training set
@@ -26,13 +35,21 @@ gdown.cached_download(
     md5="1bd99e73b2fea8e4c2ebcb0e7722f1b1",
 )
 
-# EXTRACT
+# EXTRACT training set
 torchvision.datasets.utils.extract_archive(
     from_path="CUB_train.zip",
-    to_path="Training/",
+    to_path="data/",
     remove_finished=False,
 )
 
+# CHM Weights
+gdown.cached_download(
+    url="https://drive.google.com/u/0/uc?id=1zsJRlAsoOn5F0GTCprSFYwDDfV85xDy6&export=download",
+    path="pas_psi.pt",
+    quiet=False,
+    md5="6b7b4d7bad7f89600fac340d6aa7708b",
+)
+
 
 # Caluclate Accuracy
 with open(f"./embeddings.pickle", "rb") as f:
@@ -45,35 +62,53 @@ searcher = SearchableTrainingSet(Xtrain, ytrain)
 searcher.build_index()
 
 # Extract label names
-training_folder = ImageFolder(root="./Training/train/")
+training_folder = ImageFolder(root="./data/train/")
 id_to_bird_name = {
     x[1]: x[0].split("/")[-2].replace(".", " ") for x in training_folder.imgs
 }
 
 
-def search(query_imag, searcher=searcher):
-    query_embedding = QueryToEmbedding(query_imag)
-    indices, scores, labels = searcher.search(query_embedding, k=50)
+def search(query_image, searcher=searcher):
+    query_embedding = QueryToEmbedding(query_image)
+    scores, indices, labels = searcher.search(query_embedding, k=50)
 
     result_ctr = Counter(labels[0][:20]).most_common(5)
 
     top1_label = result_ctr[0][0]
     top_indices = []
 
-    for a, b in zip(labels[0][:20], scores[0][:20]):
+    for a, b in zip(labels[0][:20], indices[0][:20]):
         if a == top1_label:
             top_indices.append(b)
 
     gallery_images = [training_folder.imgs[int(X)][0] for X in top_indices[:5]]
     predicted_labels = {id_to_bird_name[X[0]]: X[1] / 20.0 for X in result_ctr}
 
-    return predicted_labels, gallery_images
+    print("gallery_images:", gallery_images)
+
+    # CHM Prediction
+    kNN_results = (top1_label, result_ctr[0][1], gallery_images)
+    support_files = [training_folder.imgs[int(X)][0] for X in indices[0]]
+
+    print(support_files)
+    support_labels = [training_folder.imgs[int(X)][1] for X in indices[0]]
+    print(support_labels)
+
+    support = [support_files, support_labels]
+
+    chm_output = chm_classify_and_visualize(
+        query_image, kNN_results, support, training_folder
+    )
+
+    viz_plot = plot_from_reranker_output(chm_output, draw_arcs=False)
+
+    return predicted_labels, gallery_images, viz_plot
 
 
 demo = gr.Interface(
     search,
-    gr.Image(type="pil"),
-    ["label", "gallery"],
+    gr.Image(type="filepath"),
+    ["label", "gallery", "plot"],
     examples=[["./examples/bird.jpg"]],
     description="WIP - kNN on CUB dataset",
     title="Work in Progress - CHM-Corr",

common/evaluation.py

@@ -0,0 +1,32 @@
+r""" Evaluates CHMNet with PCK """
+
+import torch
+
+
+class Evaluator:
+    r""" Computes evaluation metrics of PCK """
+    @classmethod
+    def initialize(cls, alpha):
+        cls.alpha = torch.tensor(alpha).unsqueeze(1)
+
+    @classmethod
+    def evaluate(cls, prd_kps, batch):
+        r""" Compute percentage of correct key-points (PCK) with multiple alpha {0.05, 0.1, 0.15 }"""
+
+        pcks = []
+        for idx, (pk, tk) in enumerate(zip(prd_kps, batch['trg_kps'])):
+            pckthres = batch['pckthres'][idx]
+            npt = batch['n_pts'][idx]
+            prd_kps = pk[:, :npt]
+            trg_kps = tk[:, :npt]
+
+            l2dist = (prd_kps - trg_kps).pow(2).sum(dim=0).pow(0.5).unsqueeze(0).repeat(len(cls.alpha), 1)
+            thres = pckthres.expand_as(l2dist).float() * cls.alpha
+            pck = torch.le(l2dist, thres).sum(dim=1) / float(npt)
+            if len(pck) == 1: pck = pck[0]
+            pcks.append(pck)
+
+        eval_result = {'pck': pcks}
+
+        return eval_result
+

common/logger.py

@@ -0,0 +1,117 @@
+r""" Logging """
+
+import datetime
+import logging
+import os
+
+from tensorboardX import SummaryWriter
+import torch
+
+
+class Logger:
+    r""" Writes results of training/testing """
+    @classmethod
+    def initialize(cls, args, training):
+        logtime = datetime.datetime.now().__format__('_%m%d_%H%M%S')
+        logpath = args.logpath if training else '_TEST_' + args.load.split('/')[-1].split('.')[0] + logtime
+        if logpath == '': logpath = logtime
+
+        cls.logpath = os.path.join('logs', logpath + '.log')
+        cls.benchmark = args.benchmark
+        os.makedirs(cls.logpath)
+
+        logging.basicConfig(filemode='w',
+                            filename=os.path.join(cls.logpath, 'log.txt'),
+                            level=logging.INFO,
+                            format='%(message)s',
+                            datefmt='%m-%d %H:%M:%S')
+
+        # Console log config
+        console = logging.StreamHandler()
+        console.setLevel(logging.INFO)
+        formatter = logging.Formatter('%(message)s')
+        console.setFormatter(formatter)
+        logging.getLogger('').addHandler(console)
+
+        # Tensorboard writer
+        cls.tbd_writer = SummaryWriter(os.path.join(cls.logpath, 'tbd/runs'))
+
+        # Log arguments
+        if training:
+            logging.info(':======== Convolutional Hough Matching Networks =========')
+            for arg_key in args.__dict__:
+                logging.info('| %20s: %-24s' % (arg_key, str(args.__dict__[arg_key])))
+            logging.info(':========================================================\n')
+
+    @classmethod
+    def info(cls, msg):
+        r""" Writes message to .txt """
+        logging.info(msg)
+
+    @classmethod
+    def save_model(cls, model, epoch, val_pck):
+        torch.save(model.state_dict(), os.path.join(cls.logpath, 'pck_best_model.pt'))
+        cls.info('Model saved @%d w/ val. PCK: %5.2f.\n' % (epoch, val_pck))
+
+
+class AverageMeter:
+    r""" Stores loss, evaluation results, selected layers """
+    def __init__(self, benchamrk):
+        r""" Constructor of AverageMeter """
+        self.buffer_keys = ['pck']
+        self.buffer = {}
+        for key in self.buffer_keys:
+            self.buffer[key] = []
+
+        self.loss_buffer = []
+
+    def update(self, eval_result, loss=None):
+        for key in self.buffer_keys:
+            self.buffer[key] += eval_result[key]
+
+        if loss is not None:
+            self.loss_buffer.append(loss)
+
+    def write_result(self, split, epoch):
+        msg = '\n*** %s ' % split
+        msg += '[@Epoch %02d] ' % epoch
+
+        if len(self.loss_buffer) > 0:
+            msg += 'Loss: %5.2f  ' % (sum(self.loss_buffer) / len(self.loss_buffer))
+
+        for key in self.buffer_keys:
+            msg += '%s: %6.2f  ' % (key.upper(), sum(self.buffer[key]) / len(self.buffer[key]))
+        msg += '***\n'
+        Logger.info(msg)
+
+    def write_process(self, batch_idx, datalen, epoch):
+        msg = '[Epoch: %02d] ' % epoch
+        msg += '[Batch: %04d/%04d] ' % (batch_idx+1, datalen)
+        if len(self.loss_buffer) > 0:
+            msg += 'Loss: %5.2f  ' % self.loss_buffer[-1]
+            msg += 'Avg Loss: %5.5f  ' % (sum(self.loss_buffer) / len(self.loss_buffer))
+
+        for key in self.buffer_keys:
+            msg += 'Avg %s: %5.2f  ' % (key.upper(), sum(self.buffer[key]) / len(self.buffer[key]) * 100)
+        Logger.info(msg)
+
+    def write_test_process(self, batch_idx, datalen):
+        msg = '[Batch: %04d/%04d] ' % (batch_idx+1, datalen)
+
+        for key in self.buffer_keys:
+            if key == 'pck':
+                pcks = torch.stack(self.buffer[key]).mean(dim=0) * 100
+                val = ''
+                for p in pcks:
+                    val += '%5.2f   ' % p.item()
+                msg += 'Avg %s: %s   ' % (key.upper(), val)
+            else:
+                msg += 'Avg %s: %5.2f  ' % (key.upper(), sum(self.buffer[key]) / len(self.buffer[key]))
+        Logger.info(msg)
+
+    def get_test_result(self):
+        result = {}
+        for key in self.buffer_keys:
+            result[key] = torch.stack(self.buffer[key]).mean(dim=0) * 100
+
+        return result

model/base/backbone.py

@@ -0,0 +1,136 @@
+r""" ResNet-101 backbone network """
+
+import torch.utils.model_zoo as model_zoo
+import torch.nn as nn
+import torch
+
+
+__all__ = ['Backbone', 'resnet101']
+
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    r""" 3x3 convolution with padding """
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, groups=2, bias=False)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    r""" 1x1 convolution """
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, groups=2, bias=False)
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = conv1x1(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = conv1x1(planes, planes * self.expansion)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Backbone(nn.Module):
+    def __init__(self, block, layers, zero_init_residual=False):
+        super(Backbone, self).__init__()
+
+        self.inplanes = 128
+        self.conv1 = nn.Conv2d(6, 128, kernel_size=7, stride=2, padding=3, groups=2,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(128)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 128, layers[0])
+        self.layer2 = self._make_layer(block, 256, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 512, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 1024, layers[3], stride=2)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, 1000)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+
+def resnet101(pretrained=False, **kwargs):
+    """Constructs a ResNet-101 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = Backbone(Bottleneck, [3, 4, 23, 3], **kwargs)
+    if pretrained:
+        weights = model_zoo.load_url(model_urls['resnet101'])
+
+        for key in weights:
+            if key.split('.')[0] == 'fc':
+                weights[key] = weights[key].clone()
+                continue
+            weights[key] = torch.cat([weights[key].clone(), weights[key].clone()], dim=0)
+
+        model.load_state_dict(weights)
+    return model
+

model/base/chm.py

@@ -0,0 +1,190 @@
+r""" 4D and 6D convolutional Hough matching layers """
+
+from torch.nn.modules.conv import _ConvNd
+import torch.nn.functional as F
+import torch.nn as nn
+import torch
+
+from common.logger import Logger
+from . import chm_kernel
+
+
+def fast4d(corr, kernel, bias=None):
+    r""" Optimized implementation of 4D convolution """
+    bsz, ch, srch, srcw, trgh, trgw = corr.size()
+    out_channels, _, kernel_size, kernel_size, kernel_size, kernel_size = kernel.size()
+    psz = kernel_size // 2
+
+    out_corr = torch.zeros((bsz, out_channels, srch, srcw, trgh, trgw))
+    corr = corr.transpose(1, 2).contiguous().view(bsz * srch, ch, srcw, trgh, trgw)
+
+    for pidx, k3d in enumerate(kernel.permute(2, 0, 1, 3, 4, 5)):
+        inter_corr = F.conv3d(corr, k3d, bias=None, stride=1, padding=psz)
+        inter_corr = inter_corr.view(bsz, srch, out_channels, srcw, trgh, trgw).transpose(1, 2).contiguous()
+
+        add_sid = max(psz - pidx, 0)
+        add_fid = min(srch, srch + psz - pidx)
+        slc_sid = max(pidx - psz, 0)
+        slc_fid = min(srch, srch - psz + pidx)
+
+        out_corr[:, :, add_sid:add_fid, :, :, :] += inter_corr[:, :, slc_sid:slc_fid, :, :, :]
+
+    if bias is not None:
+        out_corr += bias.view(1, out_channels, 1, 1, 1, 1)
+
+    return out_corr
+
+
+def fast6d(corr, kernel, bias, diagonal_idx):
+    r""" Optimized implementation of 6D convolutional Hough matching
+         NOTE: this function only supports kernel size of (3, 3, 5, 5, 5, 5).
+    r"""
+    bsz, _, s6d, s6d, s4d, s4d, s4d, s4d = corr.size()
+    _, _, ks6d, ks6d, ks4d, ks4d, ks4d, ks4d = kernel.size()
+    corr = corr.permute(0, 2, 3, 1, 4, 5, 6, 7).contiguous().view(-1, 1, s4d, s4d, s4d, s4d)
+    kernel = kernel.view(-1, ks6d ** 2, ks4d, ks4d, ks4d, ks4d).transpose(0, 1)
+    corr = fast4d(corr, kernel).view(bsz, s6d * s6d, ks6d * ks6d, s4d, s4d, s4d, s4d)
+    corr = corr.view(bsz, s6d, s6d, ks6d, ks6d, s4d, s4d, s4d, s4d).transpose(2, 3).\
+        contiguous().view(-1, s6d * ks6d, s4d, s4d, s4d, s4d)
+
+    ndiag = s6d + (ks6d // 2) * 2
+    first_sum = []
+    for didx in diagonal_idx:
+        first_sum.append(corr[:, didx, :, :, :, :].sum(dim=1))
+    first_sum = torch.stack(first_sum).transpose(0, 1).view(bsz, s6d * ks6d, ndiag, s4d, s4d, s4d, s4d)
+
+    corr = []
+    for didx in diagonal_idx:
+        corr.append(first_sum[:, didx, :, :, :, :, :].sum(dim=1))
+    sidx = ks6d // 2
+    eidx = ndiag - sidx
+    corr = torch.stack(corr).transpose(0, 1)[:, sidx:eidx, sidx:eidx, :, :, :, :].unsqueeze(1).contiguous()
+    corr += bias.view(1, -1, 1, 1, 1, 1, 1, 1)
+
+    reverse_idx = torch.linspace(s6d * s6d - 1, 0, s6d * s6d).long()
+    corr = corr.view(bsz, 1, s6d * s6d, s4d, s4d, s4d, s4d)[:, :, reverse_idx, :, :, :, :].\
+        view(bsz, 1, s6d, s6d, s4d, s4d, s4d, s4d)
+    return corr
+
+def init_param_idx4d(param_dict):
+    param_idx = []
+    for key in param_dict:
+        curr_offset = int(key.split('_')[-1])
+        param_idx.append(torch.tensor(param_dict[key]))
+    return param_idx
+
+class CHM4d(_ConvNd):
+    r""" 4D convolutional Hough matching layer
+         NOTE: this function only supports in_channels=1 and out_channels=1.
+    r"""
+    def __init__(self, in_channels, out_channels, ksz4d, ktype, bias=True):
+        super(CHM4d, self).__init__(in_channels, out_channels, (ksz4d,) * 4,
+                                    (1,) * 4, (0,) * 4, (1,) * 4, False, (0,) * 4,
+                                    1, bias, padding_mode='zeros')
+
+        # Zero kernel initialization
+        self.zero_kernel4d = torch.zeros((in_channels, out_channels, ksz4d, ksz4d, ksz4d, ksz4d))
+        self.nkernels = in_channels * out_channels
+
+        # Initialize kernel indices
+        param_dict4d = chm_kernel.KernelGenerator(ksz4d, ktype).generate()
+        param_shared =  param_dict4d is not None
+
+        if param_shared:
+            # Initialize the shared parameters (multiplied by the number of times being shared)
+            self.param_idx = init_param_idx4d(param_dict4d)
+            weights = torch.abs(torch.randn(len(self.param_idx) * self.nkernels)) * 1e-3
+            for weight, param_idx in zip(weights.sort()[0], self.param_idx):
+                weight *= len(param_idx)
+            self.weight = nn.Parameter(weights)
+        else:  # full kernel initialziation
+            self.param_idx = None
+            self.weight = nn.Parameter(torch.abs(self.weight))
+            if bias: self.bias = nn.Parameter(torch.tensor(0.0))
+        Logger.info('(%s) # params in CHM 4D: %d' % (ktype, len(self.weight.view(-1))))
+
+    def forward(self, x):
+        kernel = self.init_kernel()
+        x = fast4d(x, kernel, self.bias)
+        return x
+
+    def init_kernel(self):
+        # Initialize CHM kernel (divided by the number of times being shared)
+        ksz = self.kernel_size[-1]
+        if self.param_idx is None:
+            kernel = self.weight
+        else:
+            kernel = torch.zeros_like(self.zero_kernel4d)
+            for idx, pdx in enumerate(self.param_idx):
+                kernel = kernel.view(-1, ksz, ksz, ksz, ksz)
+                for jdx, kernel_single in enumerate(kernel):
+                    weight = self.weight[idx + jdx * len(self.param_idx)].repeat(len(pdx)) / len(pdx)
+                    kernel_single.view(-1)[pdx] += weight
+            kernel = kernel.view(self.in_channels, self.out_channels, ksz, ksz, ksz, ksz)
+        return kernel
+
+
+class CHM6d(_ConvNd):
+    r""" 6D convolutional Hough matching layer with kernel (3, 3, 5, 5, 5, 5)
+         NOTE: this function only supports in_channels=1 and out_channels=1.
+    r"""
+    def __init__(self, in_channels, out_channels, ksz6d, ksz4d, ktype):
+        kernel_size = (ksz6d, ksz6d, ksz4d, ksz4d, ksz4d, ksz4d)
+        super(CHM6d, self).__init__(in_channels, out_channels, kernel_size, (1,) * 6,
+                                    (0,) * 6, (1,) * 6, False, (0,) * 6,
+                                    1, bias=True, padding_mode='zeros')
+
+        # Zero kernel initialization
+        self.zero_kernel4d = torch.zeros((ksz4d, ksz4d, ksz4d, ksz4d))
+        self.zero_kernel6d = torch.zeros((ksz6d, ksz6d, ksz4d, ksz4d, ksz4d, ksz4d))
+        self.nkernels = in_channels * out_channels
+
+        # Initialize kernel indices
+        # Indices in scale-space where 4D convolutions are performed (3 by 3 scale-space)
+        self.diagonal_idx = [torch.tensor(x) for x in [[6], [3, 7], [0, 4, 8], [1, 5], [2]]]
+        param_dict4d = chm_kernel.KernelGenerator(ksz4d, ktype).generate()
+        param_shared =  param_dict4d is not None
+
+        if param_shared:  # psi & iso kernel initialization
+            if ktype == 'psi':
+                self.param_dict6d = [[4], [0, 8], [2, 6], [1, 3, 5, 7]]
+            elif ktype == 'iso':
+                self.param_dict6d = [[0, 4, 8], [2, 6], [1, 3, 5, 7]]
+            self.param_dict6d = [torch.tensor(i) for i in self.param_dict6d]
+
+            # Initialize the shared parameters (multiplied by the number of times being shared)
+            self.param_idx = init_param_idx4d(param_dict4d)
+            self.param = []
+            for param_dict6d in self.param_dict6d:
+                weights = torch.abs(torch.randn(len(self.param_idx))) * 1e-3
+                for weight, param_idx in zip(weights, self.param_idx):
+                    weight *= (len(param_idx) * len(param_dict6d))
+                self.param.append(nn.Parameter(weights))
+            self.param = nn.ParameterList(self.param)
+        else:  # full kernel initialziation
+            self.param_idx = None
+            self.param = nn.Parameter(torch.abs(self.weight) * 1e-3)
+        Logger.info('(%s) # params in CHM 6D: %d' % (ktype, sum([len(x.view(-1)) for x in self.param])))
+        self.weight = None
+
+    def forward(self, corr):
+        kernel = self.init_kernel()
+        corr = fast6d(corr, kernel, self.bias, self.diagonal_idx)
+        return corr
+
+    def init_kernel(self):
+        # Initialize CHM kernel (divided by the number of times being shared)
+        if self.param_idx is None:
+            return self.param
+
+        kernel6d = torch.zeros_like(self.zero_kernel6d)
+        for idx, (param, param_dict6d) in enumerate(zip(self.param, self.param_dict6d)):
+            ksz4d = self.kernel_size[-1]
+            kernel4d = torch.zeros_like(self.zero_kernel4d)
+            for jdx, pdx in enumerate(self.param_idx):
+                kernel4d.view(-1)[pdx] += ((param[jdx] / len(pdx)) / len(param_dict6d))
+            kernel6d.view(-1, ksz4d, ksz4d, ksz4d, ksz4d)[param_dict6d] += kernel4d.view(ksz4d, ksz4d, ksz4d, ksz4d)
+        kernel6d = kernel6d.unsqueeze(0).unsqueeze(0)
+
+        return kernel6d
+

model/base/chm_kernel.py

@@ -0,0 +1,66 @@
+r""" CHM 4D kernel (psi, iso, and full) generator """
+
+import torch
+
+from .geometry import Geometry
+
+
+class KernelGenerator:
+    def __init__(self, ksz, ktype):
+        self.ksz = ksz
+        self.idx4d = Geometry.init_idx4d(ksz)
+        self.kernel = torch.zeros((ksz, ksz, ksz, ksz))
+        self.center = (ksz // 2, ksz // 2)
+        self.ktype = ktype
+
+    def quadrant(self, crd):
+        if crd[0] < self.center[0]:
+            horz_quad = -1
+        elif crd[0] < self.center[0]:
+            horz_quad = 1
+        else:
+            horz_quad = 0
+
+        if crd[1] < self.center[1]:
+            vert_quad = -1
+        elif crd[1] < self.center[1]:
+            vert_quad = 1
+        else:
+            vert_quad = 0
+
+        return horz_quad, vert_quad
+
+    def generate(self):
+        return None if self.ktype == 'full' else self.generate_chm_kernel()
+
+    def generate_chm_kernel(self):
+        param_dict = {}
+        for idx in self.idx4d:
+            src_i, src_j, trg_i, trg_j = idx
+            d_tail = Geometry.get_distance((src_i, src_j), self.center)
+            d_head = Geometry.get_distance((trg_i, trg_j), self.center)
+            d_off = Geometry.get_distance((src_i, src_j), (trg_i, trg_j))
+            horz_quad, vert_quad = self.quadrant((src_j, src_i))
+
+            src_crd = (src_i, src_j)
+            trg_crd = (trg_i, trg_j)
+
+            key = self.build_key(horz_quad, vert_quad, d_head, d_tail, src_crd, trg_crd, d_off)
+            coord1d = Geometry.get_coord1d((src_i, src_j, trg_i, trg_j), self.ksz)
+
+            if param_dict.get(key) is None: param_dict[key] = []
+            param_dict[key].append(coord1d)
+
+        return param_dict
+
+    def build_key(self, horz_quad, vert_quad, d_head, d_tail, src_crd, trg_crd, d_off):
+
+        if self.ktype == 'iso':
+            return '%d' % d_off
+        elif self.ktype == 'psi':
+            d_max = max(d_head, d_tail)
+            d_min = min(d_head, d_tail)
+            return '%d_%d_%d' % (d_max, d_min, d_off)
+        else:
+            raise Exception('not implemented.')
+

model/base/correlation.py

@@ -0,0 +1,68 @@
+r""" Provides functions that creates/manipulates correlation matrices """
+
+import math
+
+from torch.nn.functional import interpolate as resize
+import torch
+
+from .geometry import Geometry
+
+
+class Correlation:
+
+    @classmethod
+    def mutual_nn_filter(cls, correlation_matrix, eps=1e-30):
+        r""" Mutual nearest neighbor filtering (Rocco et al. NeurIPS'18 )"""
+        corr_src_max = torch.max(correlation_matrix, dim=2, keepdim=True)[0]
+        corr_trg_max = torch.max(correlation_matrix, dim=1, keepdim=True)[0]
+        corr_src_max[corr_src_max == 0] += eps
+        corr_trg_max[corr_trg_max == 0] += eps
+
+        corr_src = correlation_matrix / corr_src_max
+        corr_trg = correlation_matrix / corr_trg_max
+
+        return correlation_matrix * (corr_src * corr_trg)
+
+    @classmethod
+    def build_correlation6d(self, src_feat, trg_feat, scales, conv2ds):
+        r""" Build 6-dimensional correlation tensor """
+
+        bsz, _, side, side = src_feat.size()
+
+        # Construct feature pairs with multiple scales
+        _src_feats = []
+        _trg_feats = []
+        for scale, conv in zip(scales, conv2ds):
+            s = (round(side * math.sqrt(scale)),) * 2
+            _src_feat = conv(resize(src_feat, s, mode='bilinear', align_corners=True))
+            _trg_feat = conv(resize(trg_feat, s, mode='bilinear', align_corners=True))
+            _src_feats.append(_src_feat)
+            _trg_feats.append(_trg_feat)
+
+        # Build multiple 4-dimensional correlation tensor
+        corr6d = []
+        for src_feat in _src_feats:
+            ch = src_feat.size(1)
+
+            src_side = src_feat.size(-1)
+            src_feat = src_feat.view(bsz, ch, -1).transpose(1, 2)
+            src_norm = src_feat.norm(p=2, dim=2, keepdim=True)
+
+            for trg_feat in _trg_feats:
+                trg_side = trg_feat.size(-1)
+                trg_feat = trg_feat.view(bsz, ch, -1)
+                trg_norm = trg_feat.norm(p=2, dim=1, keepdim=True)
+
+                correlation = torch.bmm(src_feat, trg_feat) / torch.bmm(src_norm, trg_norm)
+                correlation = correlation.view(bsz, src_side, src_side, trg_side, trg_side).contiguous()
+                corr6d.append(correlation)
+
+        # Resize the spatial sizes of the 4D tensors to the same size
+        for idx, correlation in enumerate(corr6d):
+            corr6d[idx] = Geometry.interpolate4d(correlation, [side, side])
+
+        # Build 6-dimensional correlation tensor
+        corr6d = torch.stack(corr6d).view(len(scales), len(scales),
+                             bsz, side, side, side, side).permute(2, 0, 1, 3, 4, 5, 6)
+        return corr6d.clamp(min=0)
+

model/base/geometry.py

@@ -0,0 +1,133 @@
+r""" Provides functions that manipulate boxes and points """
+
+import math
+
+import torch.nn.functional as F
+import torch
+
+
+class Geometry(object):
+
+    @classmethod
+    def initialize(cls, img_size):
+        cls.img_size = img_size
+
+        cls.spatial_side = int(img_size / 8)
+        norm_grid1d = torch.linspace(-1, 1, cls.spatial_side)
+
+        cls.norm_grid_x = norm_grid1d.view(1, -1).repeat(cls.spatial_side, 1).view(1, 1, -1)
+        cls.norm_grid_y = norm_grid1d.view(-1, 1).repeat(1, cls.spatial_side).view(1, 1, -1)
+        cls.grid = torch.stack(list(reversed(torch.meshgrid(norm_grid1d, norm_grid1d)))).permute(1, 2, 0)
+
+        cls.feat_idx = torch.arange(0, cls.spatial_side).float()
+
+    @classmethod
+    def normalize_kps(cls, kps):
+        kps = kps.clone().detach()
+        kps[kps != -2] -= (cls.img_size // 2)
+        kps[kps != -2] /= (cls.img_size // 2)
+        return kps
+
+    @classmethod
+    def unnormalize_kps(cls, kps):
+        kps = kps.clone().detach()
+        kps[kps != -2] *= (cls.img_size // 2)
+        kps[kps != -2] += (cls.img_size // 2)
+        return kps
+
+    @classmethod
+    def attentive_indexing(cls, kps, thres=0.1):
+        r"""kps: normalized keypoints x, y (N, 2)
+            returns attentive index map(N, spatial_side, spatial_side)
+        """
+        nkps = kps.size(0)
+        kps = kps.view(nkps, 1, 1, 2)
+
+        eps = 1e-5
+        attmap = (cls.grid.unsqueeze(0).repeat(nkps, 1, 1, 1) - kps).pow(2).sum(dim=3)
+        attmap = (attmap + eps).pow(0.5)
+        attmap = (thres - attmap).clamp(min=0).view(nkps, -1)
+        attmap = attmap / attmap.sum(dim=1, keepdim=True)
+        attmap = attmap.view(nkps, cls.spatial_side, cls.spatial_side)
+
+        return attmap
+
+    @classmethod
+    def apply_gaussian_kernel(cls, corr, sigma=17):
+        bsz, side, side = corr.size()
+
+        center = corr.max(dim=2)[1]
+        center_y = center // cls.spatial_side
+        center_x = center % cls.spatial_side
+
+        y = cls.feat_idx.view(1, 1, cls.spatial_side).repeat(bsz, center_y.size(1), 1) - center_y.unsqueeze(2)
+        x = cls.feat_idx.view(1, 1, cls.spatial_side).repeat(bsz, center_x.size(1), 1) - center_x.unsqueeze(2)
+
+        y = y.unsqueeze(3).repeat(1, 1, 1, cls.spatial_side)
+        x = x.unsqueeze(2).repeat(1, 1, cls.spatial_side, 1)
+
+        gauss_kernel = torch.exp(-(x.pow(2) + y.pow(2)) / (2 * sigma ** 2))
+        filtered_corr = gauss_kernel * corr.view(bsz, -1, cls.spatial_side, cls.spatial_side)
+        filtered_corr = filtered_corr.view(bsz, side, side)
+
+        return filtered_corr
+
+    @classmethod
+    def transfer_kps(cls, confidence_ts, src_kps, n_pts, normalized):
+        r""" Transfer keypoints by weighted average """
+
+        if not normalized:
+            src_kps = Geometry.normalize_kps(src_kps)
+        confidence_ts = cls.apply_gaussian_kernel(confidence_ts)
+
+        pdf = F.softmax(confidence_ts, dim=2)
+        prd_x = (pdf * cls.norm_grid_x).sum(dim=2)
+        prd_y = (pdf * cls.norm_grid_y).sum(dim=2)
+
+        prd_kps = []
+        for idx, (x, y, src_kp, np) in enumerate(zip(prd_x, prd_y, src_kps, n_pts)):
+            max_pts = src_kp.size()[1]
+            prd_xy = torch.stack([x, y]).t()
+
+            src_kp = src_kp[:, :np].t()
+            attmap = cls.attentive_indexing(src_kp).view(np, -1)
+            prd_kp = (prd_xy.unsqueeze(0) * attmap.unsqueeze(-1)).sum(dim=1).t()
+            pads = (torch.zeros((2, max_pts - np)) - 2)
+            prd_kp = torch.cat([prd_kp, pads], dim=1)
+            prd_kps.append(prd_kp)
+
+        return torch.stack(prd_kps)
+
+    @staticmethod
+    def get_coord1d(coord4d, ksz):
+        i, j, k, l = coord4d
+        coord1d = i * (ksz ** 3) + j * (ksz ** 2) + k * (ksz) + l
+        return coord1d
+
+    @staticmethod
+    def get_distance(coord1, coord2):
+        delta_y = int(math.pow(coord1[0] - coord2[0], 2))
+        delta_x = int(math.pow(coord1[1] - coord2[1], 2))
+        dist = delta_y + delta_x
+        return dist
+
+    @staticmethod
+    def interpolate4d(tensor4d, size):
+        bsz, h1, w1, h2, w2 = tensor4d.size()
+        tensor4d = tensor4d.view(bsz, h1, w1, -1).permute(0, 3, 1, 2)
+        tensor4d = F.interpolate(tensor4d, size, mode='bilinear', align_corners=True)
+        tensor4d = tensor4d.view(bsz, h2, w2, -1).permute(0, 3, 1, 2)
+        tensor4d = F.interpolate(tensor4d, size, mode='bilinear', align_corners=True)
+        tensor4d = tensor4d.view(bsz, size[0], size[0], size[0], size[0])
+
+        return tensor4d
+    @staticmethod
+    def init_idx4d(ksz):
+        i0 = torch.arange(0, ksz).repeat(ksz ** 3)
+        i1 = torch.arange(0, ksz).unsqueeze(1).repeat(1, ksz).view(-1).repeat(ksz ** 2)
+        i2 = torch.arange(0, ksz).unsqueeze(1).repeat(1, ksz ** 2).view(-1).repeat(ksz)
+        i3 = torch.arange(0, ksz).unsqueeze(1).repeat(1, ksz ** 3).view(-1)
+        idx4d = torch.stack([i3, i2, i1, i0]).t().numpy()
+
+        return idx4d
+

model/chmlearner.py

@@ -0,0 +1,52 @@
+r""" Conovlutional Hough matching layers """
+
+import torch.nn as nn
+import torch
+
+from .base.correlation import Correlation
+from .base.geometry import Geometry
+from .base.chm import CHM4d, CHM6d
+
+
+class CHMLearner(nn.Module):
+
+    def __init__(self, ktype, feat_dim):
+        super(CHMLearner, self).__init__()
+
+        # Scale-wise feature transformation
+        self.scales = [0.5, 1, 2]
+        self.conv2ds = nn.ModuleList([nn.Conv2d(feat_dim, feat_dim // 4, kernel_size=3, padding=1, bias=False) for _ in self.scales])
+
+        # CHM layers
+        ksz_translation = 5
+        ksz_scale = 3
+        self.chm6d = CHM6d(1, 1, ksz_scale, ksz_translation, ktype)
+        self.chm4d = CHM4d(1, 1, ksz_translation, ktype, bias=True)
+
+        # Activations
+        self.relu = nn.ReLU(inplace=True)
+        self.sigmoid = nn.Sigmoid()
+        self.softplus = nn.Softplus()
+
+    def forward(self, src_feat, trg_feat):
+
+        corr = Correlation.build_correlation6d(src_feat, trg_feat, self.scales, self.conv2ds).unsqueeze(1)
+        bsz, ch, s, s, h, w, h, w = corr.size()
+
+        # CHM layer (6D)
+        corr = self.chm6d(corr)
+        corr = self.sigmoid(corr)
+
+        # Scale-space maxpool
+        corr = corr.view(bsz, -1, h, w, h, w).max(dim=1)[0]
+        corr = Geometry.interpolate4d(corr, [h * 2, w * 2]).unsqueeze(1)
+
+        # CHM layer (4D)
+        corr = self.chm4d(corr).squeeze(1)
+
+        # To ensure non-negative vote scores & soft cyclic constraints
+        corr = self.softplus(corr)
+        corr = Correlation.mutual_nn_filter(corr.view(bsz, corr.size(-1) ** 2, corr.size(-1) ** 2).contiguous())
+
+        return corr
+

model/chmnet.py

@@ -0,0 +1,42 @@
+r""" Convolutional Hough Matching Networks """
+
+import torch.nn as nn
+import torch
+
+from . import chmlearner as chmlearner
+from .base import backbone
+
+
+class CHMNet(nn.Module):
+    def __init__(self, ktype):
+        super(CHMNet, self).__init__()
+
+        self.backbone = backbone.resnet101(pretrained=True)
+        self.learner = chmlearner.CHMLearner(ktype, feat_dim=1024)
+
+    def forward(self, src_img, trg_img):
+        src_feat, trg_feat = self.extract_features(src_img, trg_img)
+        correlation  = self.learner(src_feat, trg_feat)
+        return correlation
+
+    def extract_features(self, src_img, trg_img):
+        feat = self.backbone.conv1.forward(torch.cat([src_img, trg_img], dim=1))
+        feat = self.backbone.bn1.forward(feat)
+        feat = self.backbone.relu.forward(feat)
+        feat = self.backbone.maxpool.forward(feat)
+
+        for idx in range(1, 5):
+            feat = self.backbone.__getattr__('layer%d' % idx)(feat)
+
+            if idx == 3:
+                src_feat = feat.narrow(1, 0, feat.size(1) // 2).clone()
+                trg_feat = feat.narrow(1, feat.size(1) // 2, feat.size(1) // 2).clone()
+                return src_feat, trg_feat
+
+    def training_objective(cls, prd_kps, trg_kps, npts):
+        l2dist = (prd_kps - trg_kps).pow(2).sum(dim=1)
+        loss = []
+        for dist, npt in zip(l2dist, npts):
+            loss.append(dist[:npt].mean())
+        return torch.stack(loss).mean()
+

visualization.py

@@ -0,0 +1,274 @@
+import pickle
+from collections import Counter
+from itertools import product
+
+import matplotlib
+import matplotlib.patches as patches
+import numpy as np
+import torchvision.transforms as transforms
+from matplotlib import gridspec
+from matplotlib import pyplot as plt
+from matplotlib.patches import ConnectionPatch, ConnectionStyle
+from PIL import Image
+
+connectionstyle = ConnectionStyle("Arc3, rad=0.2")
+
+display_transform = transforms.Compose(
+    [transforms.Resize(240), transforms.CenterCrop((240, 240))]
+)
+display_transform_knn = transforms.Compose(
+    [transforms.Resize(256), transforms.CenterCrop((224, 224))]
+)
+
+
+def keep_top_k(input_array, K=5):
+    """
+    return top 5 (k) from numpy array
+    """
+    top_5 = np.sort(input_array.reshape(-1))[::-1][K - 1]
+    masked = np.zeros_like(input_array)
+    masked[input_array >= top_5] = 1
+    return masked
+
+
+def arg_topK(inputarray, topK=5):
+    """
+    returns indicies related to top K element (largest)
+    """
+    return np.argsort(inputarray.T.reshape(-1))[::-1][:topK]
+
+
+# FOR MULTI
+def plot_from_reranker_output(reranker_output, draw_box=True, draw_arcs=True):
+    """
+    visualize chm results from a reranker output dict
+    """
+
+    ### SET COLORS
+    cmap = matplotlib.cm.get_cmap("gist_rainbow")
+    rgba = cmap(0.5)
+    colors = []
+    for k in range(5):
+        colors.append(cmap(k / 5.0))
+
+    ### SET POINTS
+    A = np.linspace(1 + 17, 240 - 17 - 1, 7)
+    point_list = list(product(A, A))
+
+    nrow = 4
+    ncol = 7
+
+    fig = plt.figure(figsize=(32, 18))
+    gs = gridspec.GridSpec(
+        nrow,
+        ncol,
+        width_ratios=[1, 0.2, 1, 1, 1, 1, 1],
+        height_ratios=[1, 1, 1, 1],
+        wspace=0.1,
+        hspace=0.1,
+        top=0.9,
+        bottom=0.05,
+        left=0.17,
+        right=0.845,
+    )
+    axes = [[None for n in range(ncol - 1)] for x in range(nrow)]
+
+    for i in range(4):
+        axes[i] = []
+        for j in range(7):
+            if j != 1:
+                if (i, j) in [(2, 0), (3, 0)]:
+                    axes[i].append(new_ax)
+                else:
+                    new_ax = plt.subplot(gs[i, j])
+                    new_ax.set_xticklabels([])
+                    new_ax.set_xticks([])
+                    new_ax.set_yticklabels([])
+                    new_ax.set_yticks([])
+                    new_ax.axis("off")
+                    axes[i].append(new_ax)
+
+    ##################### DRAW EVERYTHING
+    axes[0][0].imshow(
+        display_transform(Image.open(reranker_output["q"]).convert("RGB"))
+    )
+    axes[0][0].set_title(
+        f'Query - K={reranker_output["K"]}, N={reranker_output["N"]}', fontsize=21
+    )
+
+    axes[1][0].imshow(
+        display_transform(Image.open(reranker_output["q"]).convert("RGB"))
+    )
+    axes[1][0].set_title(f'Query - K={reranker_output["K"]}', fontsize=21)
+
+    # axes[2][0].imshow(display_transform(Image.open(reranker_output['q'])))
+
+    # CHM Top5
+    for i in range(min(5, reranker_output["chm-prediction-confidence"])):
+        axes[0][1 + i].imshow(
+            display_transform(
+                Image.open(reranker_output["chm-nearest-neighbors"][i]).convert("RGB")
+            )
+        )
+        axes[0][1 + i].set_title(f"CHM - Top - {i+1}", fontsize=21)
+
+    if reranker_output["chm-prediction-confidence"] < 5:
+        for i in range(reranker_output["chm-prediction-confidence"], 5):
+            axes[0][1 + i].imshow(Image.new(mode="RGB", size=(224, 224), color="white"))
+            axes[0][1 + i].set_title(f"", fontsize=21)
+
+    # KNN top5
+    for i in range(min(5, reranker_output["knn-prediction-confidence"])):
+        axes[1][1 + i].imshow(
+            display_transform_knn(
+                Image.open(reranker_output["knn-nearest-neighbors"][i]).convert("RGB")
+            )
+        )
+        axes[1][1 + i].set_title(f"kNN - Top - {i+1}", fontsize=21)
+
+    if reranker_output["knn-prediction-confidence"] < 5:
+        for i in range(reranker_output["knn-prediction-confidence"], 5):
+            axes[1][1 + i].imshow(Image.new(mode="RGB", size=(240, 240), color="white"))
+            axes[1][1 + i].set_title(f"", fontsize=21)
+
+    for i in range(min(5, reranker_output["chm-prediction-confidence"])):
+        axes[2][i + 1].imshow(
+            display_transform(Image.open(reranker_output["q"]).convert("RGB"))
+        )
+
+    # Lower ROWs CHM Top5
+    for i in range(min(5, reranker_output["chm-prediction-confidence"])):
+        axes[3][1 + i].imshow(
+            display_transform(
+                Image.open(reranker_output["chm-nearest-neighbors"][i]).convert("RGB")
+            )
+        )
+
+    if reranker_output["chm-prediction-confidence"] < 5:
+        for i in range(reranker_output["chm-prediction-confidence"], 5):
+            axes[2][i + 1].imshow(Image.new(mode="RGB", size=(240, 240), color="white"))
+            axes[3][1 + i].imshow(Image.new(mode="RGB", size=(240, 240), color="white"))
+
+    nzm = reranker_output["non_zero_mask"]
+    # Go throught top 5 nearest images
+
+    # #################################################################################
+    if draw_box:
+        # SQUARAES
+        for NC in range(min(5, reranker_output["chm-prediction-confidence"])):
+            # ON SOURCE
+            valid_patches_source = arg_topK(
+                reranker_output["masked_cos_values"][NC], topK=nzm
+            )
+
+            # ON QUERY
+            target_masked_patches = arg_topK(
+                reranker_output["masked_cos_values"][NC], topK=nzm
+            )
+            valid_patches_target = [
+                reranker_output["correspondance_map"][NC][x]
+                for x in target_masked_patches
+            ]
+            valid_patches_target = [(x[0] * 7) + x[1] for x in valid_patches_target]
+
+            patch_colors = [c for c in colors]
+            overlaps = [
+                item
+                for item, count in Counter(valid_patches_target).items()
+                if count > 1
+            ]
+
+            for O in overlaps:
+                indices = [i for i, val in enumerate(valid_patches_target) if val == O]
+                for ii in indices[1:]:
+                    patch_colors[ii] = patch_colors[indices[0]]
+
+            for i in valid_patches_source:
+                Psource = point_list[i]
+                rect = patches.Rectangle(
+                    (Psource[0] - 16, Psource[1] - 16),
+                    32,
+                    32,
+                    linewidth=2,
+                    edgecolor=patch_colors[valid_patches_source.tolist().index(i)],
+                    facecolor="none",
+                    alpha=1,
+                )
+                axes[2][1 + NC].add_patch(rect)
+
+            for i in valid_patches_target:
+                Psource = point_list[i]
+                rect = patches.Rectangle(
+                    (Psource[0] - 16, Psource[1] - 16),
+                    32,
+                    32,
+                    linewidth=2,
+                    edgecolor=patch_colors[valid_patches_target.index(i)],
+                    facecolor="none",
+                    alpha=1,
+                )
+                axes[3][1 + NC].add_patch(rect)
+
+    #################################################################################
+    # Show correspondence lines and points
+    if draw_arcs:
+        for CK in range(min(5, reranker_output["chm-prediction-confidence"])):
+            target_keypoints = []
+            topk_index = arg_topK(reranker_output["masked_cos_values"][CK], topK=nzm)
+            for i in range(nzm):  # Number of Connections
+                con = ConnectionPatch(
+                    xyA=(
+                        reranker_output["src-keypoints"][CK][i, 0],
+                        reranker_output["src-keypoints"][CK][i, 1],
+                    ),
+                    xyB=(
+                        reranker_output["tgt-keypoints"][CK][i, 0],
+                        reranker_output["tgt-keypoints"][CK][i, 1],
+                    ),
+                    coordsA="data",
+                    coordsB="data",
+                    axesA=axes[2][1 + CK],
+                    axesB=axes[3][1 + CK],
+                    color=colors[i],
+                    connectionstyle=connectionstyle,
+                    shrinkA=1.0,
+                    shrinkB=1.0,
+                    linewidth=1,
+                )
+
+                axes[3][1 + CK].add_artist(con)
+
+            # Scatter Plot
+            axes[2][1 + CK].scatter(
+                reranker_output["src-keypoints"][CK][:, 0],
+                reranker_output["src-keypoints"][CK][:, 1],
+                c=colors[:nzm],
+                s=10,
+            )
+            axes[3][1 + CK].scatter(
+                reranker_output["tgt-keypoints"][CK][:, 0],
+                reranker_output["tgt-keypoints"][CK][:, 1],
+                c=colors[:nzm],
+                s=10,
+            )
+
+    fig.text(
+        0.5,
+        0.95,
+        f"CHM: {reranker_output['chm-prediction']}",
+        ha="center",
+        va="bottom",
+        color="black",
+        fontsize=22,
+    )
+    fig.text(
+        0.8,
+        0.95,
+        f"KNN: {reranker_output['knn-prediction']}",
+        ha="right",
+        va="bottom",
+        color="black",
+        fontsize=22,
+    )
+
+    return fig