vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/accelerated_features/realtime_demo.py

@@ -0,0 +1,295 @@
+"""
+	"XFeat: Accelerated Features for Lightweight Image Matching, CVPR 2024."
+	https://www.verlab.dcc.ufmg.br/descriptors/xfeat_cvpr24/
+
+    Real-time homography estimation demo. Note that scene has to be planar or just rotate the camera for the estimation to work properly.
+"""
+
+import cv2
+import numpy as np
+import torch
+
+from time import time, sleep
+import argparse, sys, tqdm
+import threading
+
+from modules.xfeat import XFeat
+
+def argparser():
+    parser = argparse.ArgumentParser(description="Configurations for the real-time matching demo.")
+    parser.add_argument('--width', type=int, default=640, help='Width of the video capture stream.')
+    parser.add_argument('--height', type=int, default=480, help='Height of the video capture stream.')
+    parser.add_argument('--max_kpts', type=int, default=3_000, help='Maximum number of keypoints.')
+    parser.add_argument('--method', type=str, choices=['ORB', 'SIFT', 'XFeat'], default='XFeat', help='Local feature detection method to use.')
+    parser.add_argument('--cam', type=int, default=0, help='Webcam device number.')
+    return parser.parse_args()
+
+
+class FrameGrabber(threading.Thread):
+    def __init__(self, cap):
+        super().__init__()
+        self.cap = cap
+        _, self.frame = self.cap.read()
+        self.running = False
+
+    def run(self):
+        self.running = True
+        while self.running:
+            ret, frame = self.cap.read()
+            if not ret:
+                print("Can't receive frame (stream ended?).")
+            self.frame = frame
+            sleep(0.01)
+
+    def stop(self):
+        self.running = False
+        self.cap.release()
+
+    def get_last_frame(self):
+        return self.frame
+
+class CVWrapper():
+    def __init__(self, mtd):
+        self.mtd = mtd
+    def detectAndCompute(self, x, mask=None):
+        return self.mtd.detectAndCompute(torch.tensor(x).permute(2,0,1).float()[None])[0]
+
+class Method:
+    def __init__(self, descriptor, matcher):
+        self.descriptor = descriptor
+        self.matcher = matcher
+
+def init_method(method, max_kpts):
+    if method == "ORB":
+        return Method(descriptor=cv2.ORB_create(max_kpts, fastThreshold=10), matcher=cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True))
+    elif method == "SIFT":
+        return Method(descriptor=cv2.SIFT_create(max_kpts, contrastThreshold=-1, edgeThreshold=1000), matcher=cv2.BFMatcher(cv2.NORM_L2, crossCheck=True))
+    elif method == "XFeat":
+        return Method(descriptor=CVWrapper(XFeat(top_k = max_kpts)), matcher=XFeat())
+    else:
+        raise RuntimeError("Invalid Method.")
+
+
+class MatchingDemo:
+    def __init__(self, args):
+        self.args = args
+        self.cap = cv2.VideoCapture(args.cam)
+        self.width = args.width
+        self.height = args.height
+        self.ref_frame = None
+        self.ref_precomp = [[],[]]
+        self.corners = [[50, 50], [640-50, 50], [640-50, 480-50], [50, 480-50]]
+        self.current_frame = None
+        self.H = None
+        self.setup_camera()
+
+        #Init frame grabber thread
+        self.frame_grabber = FrameGrabber(self.cap)
+        self.frame_grabber.start()
+
+        #Homography params
+        self.min_inliers = 50
+        self.ransac_thr = 4.0
+
+        #FPS check
+        self.FPS = 0
+        self.time_list = []
+        self.max_cnt = 30 #avg FPS over this number of frames
+
+        #Set local feature method here -- we expect cv2 or Kornia convention
+        self.method = init_method(args.method, max_kpts=args.max_kpts)
+        
+        # Setting up font for captions
+        self.font = cv2.FONT_HERSHEY_SIMPLEX
+        self.font_scale = 0.9
+        self.line_type = cv2.LINE_AA
+        self.line_color = (0,255,0)
+        self.line_thickness = 3
+
+        self.window_name = "Real-time matching - Press 's' to set the reference frame."
+
+        # Removes toolbar and status bar
+        cv2.namedWindow(self.window_name, flags=cv2.WINDOW_GUI_NORMAL)
+        # Set the window size
+        cv2.resizeWindow(self.window_name, self.width*2, self.height*2)
+        #Set Mouse Callback
+        cv2.setMouseCallback(self.window_name, self.mouse_callback)
+
+    def setup_camera(self):
+        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.width)
+        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height)
+        self.cap.set(cv2.CAP_PROP_AUTO_EXPOSURE, 3)
+        #self.cap.set(cv2.CAP_PROP_EXPOSURE, 200)
+        self.cap.set(cv2.CAP_PROP_FPS, 30)
+
+        if not self.cap.isOpened():
+            print("Cannot open camera")
+            exit()
+
+    def draw_quad(self, frame, point_list):
+        if len(self.corners) > 1:
+            for i in range(len(self.corners) - 1):
+                cv2.line(frame, tuple(point_list[i]), tuple(point_list[i + 1]), self.line_color, self.line_thickness, lineType = self.line_type)
+            if len(self.corners) == 4:  # Close the quadrilateral if 4 corners are defined
+                cv2.line(frame, tuple(point_list[3]), tuple(point_list[0]), self.line_color, self.line_thickness, lineType = self.line_type)
+
+    def mouse_callback(self, event, x, y, flags, param):
+        if event == cv2.EVENT_LBUTTONDOWN:
+            if len(self.corners) >= 4:
+                self.corners = []  # Reset corners if already 4 points were clicked
+            self.corners.append((x, y))
+
+    def putText(self, canvas, text, org, fontFace, fontScale, textColor, borderColor, thickness, lineType):
+        # Draw the border
+        cv2.putText(img=canvas, text=text, org=org, fontFace=fontFace, fontScale=fontScale, 
+                    color=borderColor, thickness=thickness+2, lineType=lineType)
+        # Draw the text
+        cv2.putText(img=canvas, text=text, org=org, fontFace=fontFace, fontScale=fontScale, 
+                    color=textColor, thickness=thickness, lineType=lineType)
+
+    def warp_points(self, points, H, x_offset = 0):
+        points_np = np.array(points, dtype='float32').reshape(-1,1,2)
+
+        warped_points_np = cv2.perspectiveTransform(points_np, H).reshape(-1, 2)
+        warped_points_np[:, 0] += x_offset
+        warped_points = warped_points_np.astype(int).tolist()
+        
+        return warped_points
+
+    def create_top_frame(self):
+        top_frame_canvas = np.zeros((480, 1280, 3), dtype=np.uint8)
+        top_frame = np.hstack((self.ref_frame, self.current_frame))
+        color = (3, 186, 252)
+        cv2.rectangle(top_frame, (2, 2), (self.width*2-2, self.height-2), color, 5)  # Orange color line as a separator
+        top_frame_canvas[0:self.height, 0:self.width*2] = top_frame
+        
+        # Adding captions on the top frame canvas
+        self.putText(canvas=top_frame_canvas, text="Reference Frame:", org=(10, 30), fontFace=self.font, 
+            fontScale=self.font_scale, textColor=(0,0,0), borderColor=color, thickness=1, lineType=self.line_type)
+
+        self.putText(canvas=top_frame_canvas, text="Target Frame:", org=(650, 30), fontFace=self.font, 
+                    fontScale=self.font_scale,  textColor=(0,0,0), borderColor=color, thickness=1, lineType=self.line_type)
+        
+        self.draw_quad(top_frame_canvas, self.corners)
+        
+        return top_frame_canvas
+
+    def process(self):
+        # Create a blank canvas for the top frame
+        top_frame_canvas = self.create_top_frame()
+
+        # Match features and draw matches on the bottom frame
+        bottom_frame = self.match_and_draw(self.ref_frame, self.current_frame)
+
+        # Draw warped corners
+        if self.H is not None and len(self.corners) > 1:
+            self.draw_quad(top_frame_canvas, self.warp_points(self.corners, self.H, self.width))
+
+        # Stack top and bottom frames vertically on the final canvas
+        canvas = np.vstack((top_frame_canvas, bottom_frame))
+
+        cv2.imshow(self.window_name, canvas)
+
+    def match_and_draw(self, ref_frame, current_frame):
+
+        matches, good_matches = [], []
+        kp1, kp2 = [], []
+        points1, points2 = [], []
+
+        # Detect and compute features
+        if self.args.method in ['SIFT', 'ORB']:
+            kp1, des1 = self.ref_precomp
+            kp2, des2 = self.method.descriptor.detectAndCompute(current_frame, None)
+        else:
+            current = self.method.descriptor.detectAndCompute(current_frame)
+            kpts1, descs1 = self.ref_precomp['keypoints'], self.ref_precomp['descriptors']
+            kpts2, descs2 = current['keypoints'], current['descriptors']
+            idx0, idx1 = self.method.matcher.match(descs1, descs2, 0.82)
+            points1 = kpts1[idx0].cpu().numpy()
+            points2 = kpts2[idx1].cpu().numpy()
+
+        if len(kp1) > 10 and len(kp2) > 10 and self.args.method in ['SIFT', 'ORB']:
+            # Match descriptors
+            matches = self.method.matcher.match(des1, des2)
+
+            if len(matches) > 10:
+                points1 = np.zeros((len(matches), 2), dtype=np.float32)
+                points2 = np.zeros((len(matches), 2), dtype=np.float32)
+
+                for i, match in enumerate(matches):
+                    points1[i, :] = kp1[match.queryIdx].pt
+                    points2[i, :] = kp2[match.trainIdx].pt
+
+        if len(points1) > 10 and len(points2) > 10:
+            # Find homography
+            self.H, inliers = cv2.findHomography(points1, points2, cv2.USAC_MAGSAC, self.ransac_thr, maxIters=700, confidence=0.995)
+            inliers = inliers.flatten() > 0
+
+            if inliers.sum() < self.min_inliers:
+                self.H = None
+
+            if self.args.method in ["SIFT", "ORB"]:
+                good_matches = [m for i,m in enumerate(matches) if inliers[i]]
+            else:
+                kp1 = [cv2.KeyPoint(p[0],p[1], 5) for p in points1[inliers]]
+                kp2 = [cv2.KeyPoint(p[0],p[1], 5) for p in points2[inliers]]
+                good_matches = [cv2.DMatch(i,i,0) for i in range(len(kp1))]
+
+            # Draw matches
+            matched_frame = cv2.drawMatches(ref_frame, kp1, current_frame, kp2, good_matches, None, matchColor=(0, 200, 0), flags=2)
+            
+        else:
+            matched_frame = np.hstack([ref_frame, current_frame])
+
+        color = (240, 89, 169)
+
+        # Add a colored rectangle to separate from the top frame
+        cv2.rectangle(matched_frame, (2, 2), (self.width*2-2, self.height-2), color, 5)
+
+        # Adding captions on the top frame canvas
+        self.putText(canvas=matched_frame, text="%s Matches: %d"%(self.args.method, len(good_matches)), org=(10, 30), fontFace=self.font, 
+            fontScale=self.font_scale, textColor=(0,0,0), borderColor=color, thickness=1, lineType=self.line_type)
+        
+                # Adding captions on the top frame canvas
+        self.putText(canvas=matched_frame, text="FPS (registration): {:.1f}".format(self.FPS), org=(650, 30), fontFace=self.font, 
+            fontScale=self.font_scale, textColor=(0,0,0), borderColor=color, thickness=1, lineType=self.line_type)
+
+        return matched_frame
+
+    def main_loop(self):
+        self.current_frame = self.frame_grabber.get_last_frame()
+        self.ref_frame = self.current_frame.copy()
+        self.ref_precomp = self.method.descriptor.detectAndCompute(self.ref_frame, None) #Cache ref features
+
+        while True:
+            if self.current_frame is None:
+                break
+
+            t0 = time()
+            self.process()
+
+            key = cv2.waitKey(1)
+            if key == ord('q'):
+                break
+            elif key == ord('s'):
+                self.ref_frame = self.current_frame.copy()  # Update reference frame
+                self.ref_precomp = self.method.descriptor.detectAndCompute(self.ref_frame, None) #Cache ref features
+
+            self.current_frame = self.frame_grabber.get_last_frame()
+
+            #Measure avg. FPS
+            self.time_list.append(time()-t0)
+            if len(self.time_list) > self.max_cnt:
+                self.time_list.pop(0)
+            self.FPS = 1.0 / np.array(self.time_list).mean()
+        
+        self.cleanup()
+
+    def cleanup(self):
+        self.frame_grabber.stop()
+        self.cap.release()
+        cv2.destroyAllWindows()
+
+if __name__ == "__main__":
+    demo = MatchingDemo(args = argparser())
+    demo.main_loop()

vismatch/third_party/accelerated_features/third_party/ALIKE/alike.py

@@ -0,0 +1,143 @@
+import logging
+import os
+import cv2
+import torch
+from copy import deepcopy
+import torch.nn.functional as F
+from torchvision.transforms import ToTensor
+import math
+
+from alnet import ALNet
+from soft_detect import DKD
+import time
+
+configs = {
+    'alike-t': {'c1': 8, 'c2': 16, 'c3': 32, 'c4': 64, 'dim': 64, 'single_head': True, 'radius': 2,
+                'model_path': os.path.join(os.path.split(__file__)[0], 'models', 'alike-t.pth')},
+    'alike-s': {'c1': 8, 'c2': 16, 'c3': 48, 'c4': 96, 'dim': 96, 'single_head': True, 'radius': 2,
+                'model_path': os.path.join(os.path.split(__file__)[0], 'models', 'alike-s.pth')},
+    'alike-n': {'c1': 16, 'c2': 32, 'c3': 64, 'c4': 128, 'dim': 128, 'single_head': True, 'radius': 2,
+                'model_path': os.path.join(os.path.split(__file__)[0], 'models', 'alike-n.pth')},
+    'alike-l': {'c1': 32, 'c2': 64, 'c3': 128, 'c4': 128, 'dim': 128, 'single_head': False, 'radius': 2,
+                'model_path': os.path.join(os.path.split(__file__)[0], 'models', 'alike-l.pth')},
+}
+
+
+class ALike(ALNet):
+    def __init__(self,
+                 # ================================== feature encoder
+                 c1: int = 32, c2: int = 64, c3: int = 128, c4: int = 128, dim: int = 128,
+                 single_head: bool = False,
+                 # ================================== detect parameters
+                 radius: int = 2,
+                 top_k: int = 500, scores_th: float = 0.5,
+                 n_limit: int = 5000,
+                 device: str = 'cpu',
+                 model_path: str = ''
+                 ):
+        super().__init__(c1, c2, c3, c4, dim, single_head)
+        self.radius = radius
+        self.top_k = top_k
+        self.n_limit = n_limit
+        self.scores_th = scores_th
+        self.dkd = DKD(radius=self.radius, top_k=self.top_k,
+                       scores_th=self.scores_th, n_limit=self.n_limit)
+        self.device = device
+
+        if model_path != '':
+            state_dict = torch.load(model_path, self.device)
+            self.load_state_dict(state_dict)
+            self.to(self.device)
+            self.eval()
+            logging.info(f'Loaded model parameters from {model_path}')
+            logging.info(
+                f"Number of model parameters: {sum(p.numel() for p in self.parameters() if p.requires_grad) / 1e3}KB")
+
+    def extract_dense_map(self, image, ret_dict=False):
+        # ====================================================
+        # check image size, should be integer multiples of 2^5
+        # if it is not a integer multiples of 2^5, padding zeros
+        device = image.device
+        b, c, h, w = image.shape
+        h_ = math.ceil(h / 32) * 32 if h % 32 != 0 else h
+        w_ = math.ceil(w / 32) * 32 if w % 32 != 0 else w
+        if h_ != h:
+            h_padding = torch.zeros(b, c, h_ - h, w, device=device)
+            image = torch.cat([image, h_padding], dim=2)
+        if w_ != w:
+            w_padding = torch.zeros(b, c, h_, w_ - w, device=device)
+            image = torch.cat([image, w_padding], dim=3)
+        # ====================================================
+
+        scores_map, descriptor_map = super().forward(image)
+
+        # ====================================================
+        if h_ != h or w_ != w:
+            descriptor_map = descriptor_map[:, :, :h, :w]
+            scores_map = scores_map[:, :, :h, :w]  # Bx1xHxW
+        # ====================================================
+
+        # BxCxHxW
+        descriptor_map = torch.nn.functional.normalize(descriptor_map, p=2, dim=1)
+
+        if ret_dict:
+            return {'descriptor_map': descriptor_map, 'scores_map': scores_map, }
+        else:
+            return descriptor_map, scores_map
+
+    def forward(self, img, image_size_max=99999, sort=False, sub_pixel=False):
+        """
+        :param img: np.array HxWx3, RGB
+        :param image_size_max: maximum image size, otherwise, the image will be resized
+        :param sort: sort keypoints by scores
+        :param sub_pixel: whether to use sub-pixel accuracy
+        :return: a dictionary with 'keypoints', 'descriptors', 'scores', and 'time'
+        """
+        H, W, three = img.shape
+        assert three == 3, "input image shape should be [HxWx3]"
+
+        # ==================== image size constraint
+        image = deepcopy(img)
+        max_hw = max(H, W)
+        if max_hw > image_size_max:
+            ratio = float(image_size_max / max_hw)
+            image = cv2.resize(image, dsize=None, fx=ratio, fy=ratio)
+
+        # ==================== convert image to tensor
+        image = torch.from_numpy(image).to(self.device).to(torch.float32).permute(2, 0, 1)[None] / 255.0
+
+        # ==================== extract keypoints
+        start = time.time()
+
+        with torch.no_grad():
+            descriptor_map, scores_map = self.extract_dense_map(image)
+            keypoints, descriptors, scores, _ = self.dkd(scores_map, descriptor_map,
+                                                         sub_pixel=sub_pixel)
+            keypoints, descriptors, scores = keypoints[0], descriptors[0], scores[0]
+            keypoints = (keypoints + 1) / 2 * keypoints.new_tensor([[W - 1, H - 1]])
+
+        if sort:
+            indices = torch.argsort(scores, descending=True)
+            keypoints = keypoints[indices]
+            descriptors = descriptors[indices]
+            scores = scores[indices]
+
+        end = time.time()
+
+        return {'keypoints': keypoints.cpu().numpy(),
+                'descriptors': descriptors.cpu().numpy(),
+                'scores': scores.cpu().numpy(),
+                'scores_map': scores_map.cpu().numpy(),
+                'time': end - start, }
+
+
+if __name__ == '__main__':
+    import numpy as np
+    from thop import profile
+
+    net = ALike(c1=32, c2=64, c3=128, c4=128, dim=128, single_head=False)
+
+    image = np.random.random((640, 480, 3)).astype(np.float32)
+    flops, params = profile(net, inputs=(image, 9999, False), verbose=False)
+    print('{:<30}  {:<8} GFLops'.format('Computational complexity: ', flops / 1e9))
+    print('{:<30}  {:<8} KB'.format('Number of parameters: ', params / 1e3))

vismatch/third_party/accelerated_features/third_party/ALIKE/alnet.py

@@ -0,0 +1,164 @@
+import torch
+from torch import nn
+from torchvision.models import resnet
+from typing import Optional, Callable
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels,
+                 gate: Optional[Callable[..., nn.Module]] = None,
+                 norm_layer: Optional[Callable[..., nn.Module]] = None):
+        super().__init__()
+        if gate is None:
+            self.gate = nn.ReLU(inplace=True)
+        else:
+            self.gate = gate
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self.conv1 = resnet.conv3x3(in_channels, out_channels)
+        self.bn1 = norm_layer(out_channels)
+        self.conv2 = resnet.conv3x3(out_channels, out_channels)
+        self.bn2 = norm_layer(out_channels)
+
+    def forward(self, x):
+        x = self.gate(self.bn1(self.conv1(x)))  # B x in_channels x H x W
+        x = self.gate(self.bn2(self.conv2(x)))  # B x out_channels x H x W
+        return x
+
+
+# copied from torchvision\models\resnet.py#27->BasicBlock
+class ResBlock(nn.Module):
+    expansion: int = 1
+
+    def __init__(
+            self,
+            inplanes: int,
+            planes: int,
+            stride: int = 1,
+            downsample: Optional[nn.Module] = None,
+            groups: int = 1,
+            base_width: int = 64,
+            dilation: int = 1,
+            gate: Optional[Callable[..., nn.Module]] = None,
+            norm_layer: Optional[Callable[..., nn.Module]] = None
+    ) -> None:
+        super(ResBlock, self).__init__()
+        if gate is None:
+            self.gate = nn.ReLU(inplace=True)
+        else:
+            self.gate = gate
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError('ResBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in ResBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = resnet.conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.conv2 = resnet.conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.gate(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.gate(out)
+
+        return out
+
+
+class ALNet(nn.Module):
+    def __init__(self, c1: int = 32, c2: int = 64, c3: int = 128, c4: int = 128, dim: int = 128,
+                 single_head: bool = True,
+                 ):
+        super().__init__()
+
+        self.gate = nn.ReLU(inplace=True)
+
+        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.pool4 = nn.MaxPool2d(kernel_size=4, stride=4)
+
+        self.block1 = ConvBlock(3, c1, self.gate, nn.BatchNorm2d)
+
+        self.block2 = ResBlock(inplanes=c1, planes=c2, stride=1,
+                               downsample=nn.Conv2d(c1, c2, 1),
+                               gate=self.gate,
+                               norm_layer=nn.BatchNorm2d)
+        self.block3 = ResBlock(inplanes=c2, planes=c3, stride=1,
+                               downsample=nn.Conv2d(c2, c3, 1),
+                               gate=self.gate,
+                               norm_layer=nn.BatchNorm2d)
+        self.block4 = ResBlock(inplanes=c3, planes=c4, stride=1,
+                               downsample=nn.Conv2d(c3, c4, 1),
+                               gate=self.gate,
+                               norm_layer=nn.BatchNorm2d)
+
+        # ================================== feature aggregation
+        self.conv1 = resnet.conv1x1(c1, dim // 4)
+        self.conv2 = resnet.conv1x1(c2, dim // 4)
+        self.conv3 = resnet.conv1x1(c3, dim // 4)
+        self.conv4 = resnet.conv1x1(dim, dim // 4)
+        self.upsample2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.upsample4 = nn.Upsample(scale_factor=4, mode='bilinear', align_corners=True)
+        self.upsample8 = nn.Upsample(scale_factor=8, mode='bilinear', align_corners=True)
+        self.upsample32 = nn.Upsample(scale_factor=32, mode='bilinear', align_corners=True)
+
+        # ================================== detector and descriptor head
+        self.single_head = single_head
+        if not self.single_head:
+            self.convhead1 = resnet.conv1x1(dim, dim)
+        self.convhead2 = resnet.conv1x1(dim, dim + 1)
+
+    def forward(self, image):
+        # ================================== feature encoder
+        x1 = self.block1(image)  # B x c1 x H x W
+        x2 = self.pool2(x1)
+        x2 = self.block2(x2)  # B x c2 x H/2 x W/2
+        x3 = self.pool4(x2)
+        x3 = self.block3(x3)  # B x c3 x H/8 x W/8
+        x4 = self.pool4(x3)
+        x4 = self.block4(x4)  # B x dim x H/32 x W/32
+
+        # ================================== feature aggregation
+        x1 = self.gate(self.conv1(x1))  # B x dim//4 x H x W
+        x2 = self.gate(self.conv2(x2))  # B x dim//4 x H//2 x W//2
+        x3 = self.gate(self.conv3(x3))  # B x dim//4 x H//8 x W//8
+        x4 = self.gate(self.conv4(x4))  # B x dim//4 x H//32 x W//32
+        x2_up = self.upsample2(x2)  # B x dim//4 x H x W
+        x3_up = self.upsample8(x3)  # B x dim//4 x H x W
+        x4_up = self.upsample32(x4)  # B x dim//4 x H x W
+        x1234 = torch.cat([x1, x2_up, x3_up, x4_up], dim=1)
+
+        # ================================== detector and descriptor head
+        if not self.single_head:
+            x1234 = self.gate(self.convhead1(x1234))
+        x = self.convhead2(x1234)  # B x dim+1 x H x W
+
+        descriptor_map = x[:, :-1, :, :]
+        scores_map = torch.sigmoid(x[:, -1, :, :]).unsqueeze(1)
+
+        return scores_map, descriptor_map
+
+
+if __name__ == '__main__':
+    from thop import profile
+
+    net = ALNet(c1=16, c2=32, c3=64, c4=128, dim=128, single_head=True)
+
+    image = torch.randn(1, 3, 640, 480)
+    flops, params = profile(net, inputs=(image,), verbose=False)
+    print('{:<30}  {:<8} GFLops'.format('Computational complexity: ', flops / 1e9))
+    print('{:<30}  {:<8} KB'.format('Number of parameters: ', params / 1e3))