vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/TopicFM/viz/methods/base.py

@@ -0,0 +1,70 @@
+import pprint
+from abc import ABCMeta, abstractmethod
+import torch
+from itertools import chain
+import cv2
+
+from src.utils.plotting import make_matching_figure, error_colormap
+from src.utils.metrics import aggregate_metrics
+
+
+def flatten_list(x):
+    return list(chain(*x))
+
+
+class Viz(metaclass=ABCMeta):
+    def __init__(self):
+        super().__init__()
+        self.device = torch.device('cuda:{}'.format(0) if torch.cuda.is_available() else 'cpu')
+        torch.set_grad_enabled(False)
+
+        # for evaluation metrics of MegaDepth and ScanNet
+        self.eval_stats = []
+        self.time_stats = []
+        self.flops_stats = {"backbone": [], "coarse_net": [], "fine_net": [], "total": []}
+
+    def draw_matches(self, mkpts0, mkpts1, img0, img1, conf, path=None, **kwargs):
+        thr = 5e-4
+        img0 = cv2.cvtColor(img0, cv2.COLOR_BGR2RGB)
+        img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
+        if "conf_thr" in kwargs:
+            thr = kwargs["conf_thr"]
+        color = error_colormap(conf, thr, alpha=0.1)
+
+        text = [
+            f"{self.name}",
+            f"#Matches: {len(mkpts0)}",
+        ]
+        if 'R_errs' in kwargs:
+            text.append(f"$\\Delta$R:{kwargs['R_errs']:.2f}°,  $\\Delta$t:{kwargs['t_errs']:.2f}°",)
+
+        if path:
+            make_matching_figure(img0, img1, mkpts0, mkpts1, color, text=text, path=path, dpi=150)
+        else:
+            return make_matching_figure(img0, img1, mkpts0, mkpts1, color, text=text)
+
+    @abstractmethod
+    def match_and_draw(self, data_dict, **kwargs):
+        pass
+
+    def compute_eval_metrics(self, epi_err_thr=5e-4):
+        # metrics: dict of list, numpy
+        _metrics = [o['metrics'] for o in self.eval_stats]
+        metrics = {k: flatten_list([_me[k] for _me in _metrics]) for k in _metrics[0]}
+
+        val_metrics_4tb = aggregate_metrics(metrics, epi_err_thr)
+        print('\n' + pprint.pformat(val_metrics_4tb))
+
+    def measure_time(self):
+        if len(self.time_stats) == 0:
+            return 0
+        return sum(self.time_stats) / len(self.time_stats)
+
+    def measure_flops(self):
+        outputs = {}
+        for k, v in self.flops_stats.items():
+            if len(v) == 0:
+                outputs[k] = 0
+            else:
+                outputs[k] = sum(v) / len(v)
+        return outputs

vismatch/third_party/TopicFM/viz/methods/topicfmv2.py

@@ -0,0 +1,208 @@
+from argparse import Namespace
+import os
+import torch
+import cv2
+from time import time
+from pathlib import Path
+import matplotlib.cm as cm
+import numpy as np
+
+from src.models.topic_fm import TopicFM
+from src import get_model_cfg
+from .base import Viz
+from src.utils.metrics import compute_symmetrical_epipolar_errors, compute_pose_errors
+from src.utils.plotting import draw_topics, draw_topicfm_demo, error_colormap
+
+
+class VizTopicFMv2(Viz):
+    def __init__(self, args):
+        super().__init__()
+        if type(args) == dict:
+            args = Namespace(**args)
+
+        self.match_threshold = args.match_threshold
+        # self.n_sampling_topics = args.n_sampling_topics
+        self.show_n_topics = args.show_n_topics
+
+        # Load model
+        conf = dict(get_model_cfg())
+        conf['match_coarse']['thr'] = self.match_threshold
+        # conf['coarse']['n_samples'] = self.n_sampling_topics
+        for k, v in args.coarse_model_cfg.items():
+            conf["coarse"][k] = v
+        conf['loss']['fine_type'] = "sym_epi"
+        print("model config: ", conf)
+        self.model = TopicFM(config=conf)
+        ckpt_dict = torch.load(args.ckpt)
+        self.model.load_state_dict(ckpt_dict['state_dict'])
+        self.model = self.model.eval().to(self.device)
+
+        # Name the method
+        # self.ckpt_name = args.ckpt.split('/')[-1].split('.')[0]
+        self.name = 'TopicFMv2'
+
+        print(f'Initialize {self.name}')
+
+    def match_and_draw(self, data_dict, root_dir=None, ground_truth=False, measure_time=False, viz_matches=True):
+        if measure_time:
+            torch.cuda.synchronize()
+            start = torch.cuda.Event(enable_timing=True)
+            end = torch.cuda.Event(enable_timing=True)
+            start.record()
+        self.model(data_dict)
+        if measure_time:
+            torch.cuda.synchronize()
+            end.record()
+            torch.cuda.synchronize()
+            self.time_stats.append(start.elapsed_time(end))
+
+            # self.flops_stats["backbone"].append(data_dict["backbone_flops"])
+            # self.flops_stats["coarse_net"].append(data_dict["coarse_net_flops"])
+            # self.flops_stats["fine_net"].append(data_dict["fine_net_flops"])
+            # self.flops_stats["total"].append(data_dict["backbone_flops"] + data_dict["coarse_net_flops"] +
+            #                                  data_dict["fine_net_flops"])
+
+
+        kpts0 = data_dict['mkpts0_f'].cpu().numpy()
+        kpts1 = data_dict['mkpts1_f'].cpu().numpy()
+
+        img_name0, img_name1 = list(zip(*data_dict['pair_names']))[0]
+        img0 = cv2.imread(os.path.join(root_dir, img_name0))
+        img1 = cv2.imread(os.path.join(root_dir, img_name1))
+        if str(data_dict["dataset_name"][0]).lower() == 'scannet':
+            img0 = cv2.resize(img0, (640, 480))
+            img1 = cv2.resize(img1, (640, 480))
+
+        if viz_matches:
+            saved_name = "_".join([img_name0.split('/')[-1].split('.')[0], img_name1.split('/')[-1].split('.')[0]])
+            folder_matches = os.path.join(root_dir, "{}_viz_matches".format(self.name))
+            if not os.path.exists(folder_matches):
+                os.makedirs(folder_matches)
+            path_to_save_matches = os.path.join(folder_matches, "{}.png".format(saved_name))
+
+            if ground_truth:
+                compute_symmetrical_epipolar_errors(data_dict)  # compute epi_errs for each match
+                compute_pose_errors(data_dict, ransac_thr=0.4)  # compute R_errs, t_errs, pose_errs for each pair
+                epi_errors = data_dict['epi_errs'].cpu().numpy()
+                R_errors, t_errors = data_dict['R_errs'][0], data_dict['t_errs'][0]
+
+                # self.draw_matches(kpts0, kpts1, img0, img1, epi_errors, path=path_to_save_matches,
+                #                   R_errs=R_errors, t_errs=t_errors)
+
+                # compute evaluation metrics
+                rel_pair_names = list(zip(*data_dict['pair_names']))
+                bs = data_dict['image0'].size(0)
+                metrics = {
+                    # to filter duplicate pairs caused by DistributedSampler
+                    'identifiers': ['#'.join(rel_pair_names[b]) for b in range(bs)],
+                    'epi_errs': [data_dict['epi_errs'][data_dict['m_bids'] == b].cpu().numpy() for b in range(bs)],
+                    'R_errs': data_dict['R_errs'],
+                    't_errs': data_dict['t_errs'],
+                    'inliers': data_dict['inliers']}
+                self.eval_stats.append({'metrics': metrics})
+            else:
+                m_conf = 1 - data_dict["mconf"].cpu().numpy()
+                self.draw_matches(kpts0, kpts1, img0, img1, m_conf, path=path_to_save_matches, conf_thr=0.4)
+            if self.show_n_topics > 0:
+                folder_topics = os.path.join(root_dir, "{}_viz_topics".format(self.name))
+                if not os.path.exists(folder_topics):
+                    os.makedirs(folder_topics)
+                draw_topics(data_dict, img0, img1, saved_folder=folder_topics, show_n_topics=self.show_n_topics,
+                            saved_name=saved_name)
+
+    def run_demo(self, dataloader, writer=None, output_dir=None, no_display=False, skip_frames=1):
+        data_dict = next(dataloader)
+
+        frame_id = 0
+        last_image_id = 0
+        img0 = np.array(cv2.imread(str(data_dict["img_path"][0])), dtype=np.float32) / 255
+        frame_tensor = data_dict["img"].to(self.device)
+        pair_data = {'image0': frame_tensor}
+        last_frame = cv2.resize(img0, (frame_tensor.shape[-1], frame_tensor.shape[-2]), cv2.INTER_LINEAR)
+
+        if output_dir is not None:
+            print('==> Will write outputs to {}'.format(output_dir))
+            Path(output_dir).mkdir(exist_ok=True)
+
+        # Create a window to display the demo.
+        if not no_display:
+            window_name = 'Topic-assisted Feature Matching'
+            cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
+            cv2.resizeWindow(window_name, (640 * 2, 480 * 2))
+        else:
+            print('Skipping visualization, will not show a GUI.')
+
+        # Print the keyboard help menu.
+        print('==> Keyboard control:\n'
+              '\tn: select the current frame as the reference image (left)\n'
+              '\tq: quit')
+
+        # vis_range = [kwargs["bottom_k"], kwargs["top_k"]]
+
+        while True:
+            frame_id += 1
+            if frame_id == len(dataloader):
+                print('Finished demo_loftr.py')
+                break
+            data_dict = next(dataloader)
+            if frame_id % skip_frames != 0:
+                # print("Skipping frame.")
+                continue
+
+            stem0, stem1 = last_image_id, data_dict["id"][0].item() - 1
+            frame = np.array(cv2.imread(str(data_dict["img_path"][0])), dtype=np.float32) / 255
+
+            frame_tensor = data_dict["img"].to(self.device)
+            frame = cv2.resize(frame, (frame_tensor.shape[-1], frame_tensor.shape[-2]), interpolation=cv2.INTER_LINEAR)
+            pair_data = {**pair_data, 'image1': frame_tensor}
+            self.model(pair_data)
+
+            total_n_matches = len(pair_data['mkpts0_f'])
+            mkpts0 = pair_data['mkpts0_f'].cpu().numpy()  # [vis_range[0]:vis_range[1]]
+            mkpts1 = pair_data['mkpts1_f'].cpu().numpy()  # [vis_range[0]:vis_range[1]]
+            mconf = pair_data['mconf'].cpu().numpy()  # [vis_range[0]:vis_range[1]]
+
+            # Normalize confidence.
+            if len(mconf) > 0:
+                mconf = 1 - mconf
+
+            # alpha = 0
+            # color = cm.jet(mconf, alpha=alpha)
+            color = error_colormap(mconf, thr=0.4, alpha=0.1)
+
+            text = [
+                f'Topics',
+                '#Matches: {}'.format(total_n_matches),
+            ]
+
+            out = draw_topicfm_demo(pair_data, last_frame, frame, mkpts0, mkpts1, color, text,
+                                    show_n_topics=4, path=None)
+
+            if not no_display:
+                if writer is not None:
+                    writer.write(out)
+                cv2.imshow('TopicFM Matches', out)
+                key = chr(cv2.waitKey(10) & 0xFF)
+                if key == 'q':
+                    if writer is not None:
+                        writer.release()
+                    print('Exiting...')
+                    break
+                elif key == 'n':
+                    pair_data['image0'] = frame_tensor
+                    last_frame = frame
+                    last_image_id = (data_dict["id"][0].item() - 1)
+                    frame_id_left = frame_id
+
+            elif output_dir is not None:
+                stem = 'matches_{:06}_{:06}'.format(stem0, stem1)
+                out_file = str(Path(output_dir, stem + '.png'))
+                print('\nWriting image to {}'.format(out_file))
+                cv2.imwrite(out_file, out)
+            else:
+                raise ValueError("output_dir is required when no display is given.")
+
+        cv2.destroyAllWindows()
+        if writer is not None:
+            writer.release()
+

vismatch/third_party/UFM/UniCeption/examples/models/cosmos/autoencoding.py

@@ -0,0 +1,48 @@
+import os
+
+import cv2
+import torch
+from matplotlib import pyplot as plt
+
+from uniception.models.encoders.base import ViTEncoderInput
+from uniception.models.encoders.cosmos import CosmosEncoder
+from uniception.models.prediction_heads.cosmos import CosmosSingleChannel
+
+base_path = os.path.dirname(os.path.abspath(__file__))
+
+encoder = CosmosEncoder(
+    name="cosmos",
+    patch_size=8,
+    pretrained_checkpoint_path=os.path.join(
+        base_path, "../../../checkpoints/encoders/cosmos/Cosmos-Tokenizer-CI8x8/encoder.pth"
+    ),
+)
+
+decoder = CosmosSingleChannel(
+    patch_size=8,
+    pretrained_checkpoint_path=os.path.join(base_path, "../../../checkpoints/prediction_heads/cosmos/decoder_8.pth"),
+)
+
+example_image = cv2.imread(os.path.join(base_path, "./example.png"))
+example_image = cv2.cvtColor(example_image, cv2.COLOR_BGR2RGB)
+example_tensor = torch.tensor(example_image).permute(2, 0, 1).unsqueeze(0).float() / 255.0
+example_tensor = example_tensor * 2.0 - 1.0  # Normalize to [-1, 1] according to the COSMOS Encoder
+
+encoded_latent = encoder(ViTEncoderInput("cosmos", example_tensor)).features
+
+decoded_image = decoder(encoded_latent)
+decoded_image = (decoded_image + 1.0) / 2.0  # Denormalize to [0, 1] for visualization
+
+# plot the original and decoded images
+plt.figure(figsize=(10, 5))
+plt.subplot(1, 2, 1)
+plt.imshow(example_image)
+plt.title("Original Image")
+plt.axis("off")
+
+plt.subplot(1, 2, 2)
+plt.imshow(decoded_image.squeeze().detach().permute(1, 2, 0).cpu().numpy())
+plt.title("Decoded Image")
+plt.axis("off")
+
+plt.savefig(os.path.join(base_path, "example_decoded.png"))

vismatch/third_party/UFM/UniCeption/examples/models/dust3r/convert_dust3r_weights_to_uniception.py

@@ -0,0 +1,331 @@
+"""
+This file extracts the cross-attention transformer & prediction head weights from dust3r checkpoints into uniception format.
+
+Special Notice: dust3r have changed their released weights before/after CVPR, and
+uniception uses the checkpoint BEFORE CVPR (they perform better). So please make sure you are not converting
+the newly downloaded weights. Consult Yuchen and Nikhil on where to find the old weights.
+"""
+
+import argparse
+import os
+
+import torch
+from torch import nn
+
+from uniception.models.info_sharing.cross_attention_transformer import MultiViewCrossAttentionTransformerIFR
+from uniception.models.prediction_heads.dpt import DPTFeature, DPTRegressionProcessor
+from uniception.models.prediction_heads.linear import LinearFeature
+
+
+def extract_cross_attention_weights(checkpoint_path, output_folder, output_filename):
+    "Extract the UniCeption format cross attention weights from the original CroCoV2/DUSt3R/MASt3R checkpoints."
+    # Load checkpoint
+    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+
+    # Filter the relevant keys for the cross attention model and duplicate if necessary
+    filtered_checkpoint = checkpoint["model"]
+    filtered_checkpoint = {k: v for k, v in filtered_checkpoint.items() if "dec" in k}
+    duplicate_checkpoint = {}
+    if not any(k.startswith("dec_blocks2") for k in filtered_checkpoint):
+        print("Duplicating dec_blocks to dec_blocks2")
+        for key, value in filtered_checkpoint.items():
+            if key.startswith("dec_blocks"):
+                duplicate_checkpoint[key.replace("dec_blocks", "dec_blocks2")] = value
+        filtered_checkpoint = {**filtered_checkpoint, **duplicate_checkpoint}
+    new_checkpoint = {}
+    for k, v in filtered_checkpoint.items():
+        if "decoder_embed" in k:
+            new_key = k.replace("decoder_embed", "proj_embed")
+            new_checkpoint[new_key] = v
+        elif "dec_blocks." in k:
+            new_key = k.replace("dec_blocks.", "multi_view_branches.0.")
+            new_checkpoint[new_key] = v
+        elif "dec_blocks2." in k:
+            new_key = k.replace("dec_blocks2.", "multi_view_branches.1.")
+            new_checkpoint[new_key] = v
+        elif "dec_norm" in k:
+            new_key = k.replace("dec_norm", "norm")
+            new_checkpoint[new_key] = v
+
+    # Init model
+    model = MultiViewCrossAttentionTransformerIFR(
+        name="MV-CAT-IFR",
+        input_embed_dim=1024,
+        num_views=2,
+        indices=[5, 8],
+        norm_intermediate=False,
+    )
+
+    # Load new checkpoint
+    print(model.load_state_dict(new_checkpoint))
+
+    # Save the checkpoint
+    save_checkpoint = {}
+    save_checkpoint["model"] = model.state_dict()
+    os.makedirs(os.path.join(output_folder, "cross_attn_transformer"), exist_ok=True)
+    save_path = os.path.join(output_folder, "cross_attn_transformer", output_filename)
+    torch.save(save_checkpoint, save_path)
+
+
+def extract_dust3r_dpt_checkpoints(checkpoint_path, output_folder, output_filename):
+    "Extract the UniCeption format DPT head weights from the original DUSt3R checkpoint."
+    source_ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+
+    for head in ["head1", "head2"]:
+        # Extract head weights from the checkpoint
+        dpt_head_weights = {k: v for k, v in source_ckpt["model"].items() if k.startswith(f"downstream_{head}")}
+        dpt_head_weights = {k.replace(f"downstream_{head}.dpt.", ""): v for k, v in dpt_head_weights.items()}
+        dpt_feature_weights = {k: v for k, v in dpt_head_weights.items() if not (k.startswith("head"))}
+
+        # Construct the DPTFeature module and load the weights
+        dpt = DPTFeature(
+            patch_size=16,
+            hooks=[0, 1, 2, 3],
+            input_feature_dims=[1024, 768, 768, 768],
+            layer_dims=[96, 192, 384, 768],
+            feature_dim=256,
+            use_bn=False,
+            output_width_ratio=1,
+        )
+
+        dpt.load_state_dict(dpt_feature_weights, strict=True)
+
+        # Construct the dpt processor module and load the weights
+        dpt_processor_weights = {k.replace("head.", ""): v for k, v in dpt_head_weights.items() if k.startswith("head")}
+
+        # Replace the keys according to:
+        key_replace_dict = {
+            "0.weight": "conv1.weight",
+            "0.bias": "conv1.bias",
+            "2.weight": "conv2.0.weight",
+            "2.bias": "conv2.0.bias",
+            "4.weight": "conv2.2.weight",
+            "4.bias": "conv2.2.bias",
+        }
+
+        dpt_processor_weights = {key_replace_dict.get(k, k): v for k, v in dpt_processor_weights.items()}
+
+        dpt_reg_processor = DPTRegressionProcessor(input_feature_dim=256, output_dim=4, hidden_dims=[128, 128])
+
+        dpt_reg_processor.load_state_dict(dpt_processor_weights, strict=True)
+
+        # Save the state_dicts of the DPTFeature and DPTRegressionProcessor
+        dpt_feature_path = os.path.join(output_folder, "dpt_feature_head", output_filename + f"_feature_{head}.pth")
+        dpt_reg_processor_path = os.path.join(
+            output_folder, "dpt_reg_processor", output_filename + f"_reg_processor{head[-1]}.pth"
+        )
+
+        os.makedirs(os.path.dirname(dpt_feature_path), exist_ok=True)
+        os.makedirs(os.path.dirname(dpt_reg_processor_path), exist_ok=True)
+
+        torch.save({"model": dpt.state_dict()}, dpt_feature_path)
+        torch.save({"model": dpt_reg_processor.state_dict()}, dpt_reg_processor_path)
+
+
+def extract_dust3r_linear_checkpoints(checkpoint_path, output_folder, output_filename):
+    "Extract the UniCeption format linear head weights from the original DUSt3R checkpoint."
+    test_linear_to_conv()
+
+    source_ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+
+    for head in ["head1", "head2"]:
+        linear_head_params = {k: v for k, v in source_ckpt["model"].items() if k.startswith(f"downstream_{head}")}
+        linear_head_params = {k.replace(f"downstream_{head}.proj.", ""): v for k, v in linear_head_params.items()}
+
+        assert set(linear_head_params.keys()) == {"weight", "bias"}
+
+        input_feature_dim = 768
+        output_dim = 4
+        patch_size = 16
+
+        linear = nn.Linear(input_feature_dim, output_dim * patch_size * patch_size, bias=True)
+        linear.load_state_dict(linear_head_params, strict=True)
+
+        conv_layer = linear_to_conv2d(linear)
+
+        linear_feature = LinearFeature(input_feature_dim, 4, patch_size)
+        linear_feature.linear.load_state_dict(conv_layer.state_dict(), strict=True)
+
+        linear_feature_path = os.path.join(
+            output_folder, "linear_feature_head", output_filename + f"_feature_{head}.pth"
+        )
+        os.makedirs(os.path.dirname(linear_feature_path), exist_ok=True)
+        torch.save({"model": linear_feature.state_dict()}, linear_feature_path)
+
+
+def extract_mast3r_dpt_checkpoints(checkpoint_path, output_folder, output_filename):
+    "Extract the UniCeption format DPT head weights from the original MASt3R checkpoint."
+    source_ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+
+    for head in ["head1", "head2"]:
+        dpt_head = {k: v for k, v in source_ckpt["model"].items() if k.startswith(f"downstream_{head}")}
+        dpt_head = {k.replace(f"downstream_{head}.", ""): v for k, v in dpt_head.items()}
+        dpt_head = {k.replace("dpt.", ""): v for k, v in dpt_head.items()}
+
+        dpt_feature_weights = {
+            k: v for k, v in dpt_head.items() if not (k.startswith("head") or k.startswith("head_local_features"))
+        }
+
+        dpt = DPTFeature(
+            patch_size=16,
+            hooks=[0, 1, 2, 3],
+            input_feature_dims=[1024, 768, 768, 768],
+            layer_dims=[96, 192, 384, 768],
+            feature_dim=256,
+            use_bn=False,
+            output_width_ratio=1,
+        )
+
+        dpt.load_state_dict(dpt_feature_weights, strict=True)
+
+        dpt_processor_weights = {
+            k.replace("head.", ""): v
+            for k, v in dpt_head.items()
+            if (k.startswith("head") and not k.startswith("head_local_features"))
+        }
+
+        # Replace the keys according to:
+        key_replace_dict = {
+            "0.weight": "conv1.weight",
+            "0.bias": "conv1.bias",
+            "2.weight": "conv2.0.weight",
+            "2.bias": "conv2.0.bias",
+            "4.weight": "conv2.2.weight",
+            "4.bias": "conv2.2.bias",
+        }
+
+        dpt_processor_weights = {key_replace_dict.get(k, k): v for k, v in dpt_processor_weights.items()}
+
+        dpt_reg_processor = DPTRegressionProcessor(input_feature_dim=256, output_dim=4, hidden_dims=[128, 128])
+
+        dpt_reg_processor.load_state_dict(dpt_processor_weights, strict=True)
+
+        # Save the state_dicts of the DPTFeature and DPTRegressionProcessor
+        dpt_feature_path = os.path.join(output_folder, "dpt_feature_head", output_filename + f"_feature_{head}.pth")
+        dpt_reg_processor_path = os.path.join(
+            output_folder, "dpt_reg_processor", output_filename + f"_reg_processor{head[-1]}.pth"
+        )
+
+        os.makedirs(os.path.dirname(dpt_feature_path), exist_ok=True)
+        os.makedirs(os.path.dirname(dpt_reg_processor_path), exist_ok=True)
+
+        torch.save({"model": dpt.state_dict()}, dpt_feature_path)
+        torch.save({"model": dpt_reg_processor.state_dict()}, dpt_reg_processor_path)
+
+
+def linear_to_conv2d(linear_layer):
+    """
+    Converts a nn.Linear layer to an equivalent nn.Conv2d layer with a 1x1 kernel.
+
+    Parameters:
+    - linear_layer (nn.Linear): The Linear layer to convert.
+
+    Returns:
+    - conv_layer (nn.Conv2d): The equivalent Conv2d layer.
+    """
+    # Extract in_features and out_features from the Linear layer
+    in_features = linear_layer.in_features
+    out_features = linear_layer.out_features
+    bias = linear_layer.bias is not None
+
+    # Create a Conv2d layer with a 1x1 kernel
+    conv_layer = nn.Conv2d(
+        in_channels=in_features, out_channels=out_features, kernel_size=1, stride=1, padding=0, bias=bias
+    )
+
+    # Reshape Linear weights to match Conv2d weights
+    conv_weight = linear_layer.weight.data.view(out_features, in_features, 1, 1).clone()
+    conv_layer.weight.data = conv_weight
+
+    # Copy bias if it exists
+    if bias:
+        conv_layer.bias.data = linear_layer.bias.data.clone()
+
+    return conv_layer
+
+
+def test_linear_to_conv():
+    "Test the linear_to_conv2d function."
+    batch_size = 4
+    height = 16
+    width = 24
+    in_channels = 3
+    out_channels = 5
+
+    # Sample input tensor in BHWC format
+    x_linear = torch.randn(batch_size, height, width, in_channels)
+
+    # Define Linear layer
+    linear_layer = nn.Linear(in_channels, out_channels)
+    output_linear = linear_layer(x_linear)
+
+    # Transpose input tensor to BCHW format for Conv2d
+    x_conv = x_linear.permute(0, 3, 1, 2)
+
+    # Define Conv2d layer
+    conv_layer = linear_to_conv2d(linear_layer)
+
+    # Get Conv2d output and transpose back to BHWC format
+    output_conv = conv_layer(x_conv).permute(0, 2, 3, 1)
+
+    # Verify that outputs are the same
+    assert torch.allclose(output_linear, output_conv, atol=1e-6)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Extract dust3r checkpoints to uniception format")
+
+    parser.add_argument(
+        "-dcf", "--dust3r_checkpoints_folder", type=str, required=True, help="Path to the dust3r checkpoints folder"
+    )
+    parser.add_argument("-of", "--output_folder", type=str, required=True, help="Path to the output folder")
+
+    args = parser.parse_args()
+
+    output_folder = args.output_folder
+    info_sharing_output_folder = os.path.join(output_folder, "info_sharing")
+    pred_head_output_folder = os.path.join(output_folder, "prediction_heads")
+    os.makedirs(output_folder, exist_ok=True)
+    os.makedirs(info_sharing_output_folder, exist_ok=True)
+    os.makedirs(pred_head_output_folder, exist_ok=True)
+
+    # Extract croco checkpoint
+    print("Extracting CroCo checkpoint...")
+    croco_ckpt_filepath = os.path.join(args.dust3r_checkpoints_folder, "CroCo_V2_ViTLarge_BaseDecoder.pth")
+    extract_cross_attention_weights(
+        croco_ckpt_filepath, info_sharing_output_folder, "Two_View_Cross_Attention_Transformer_CroCo.pth"
+    )
+
+    # Extract dust3r 224 linear checkpoint
+    print("Extracting DUSt3R 224 linear checkpoint...")
+    dust3r_ckpt_filepath = os.path.join(args.dust3r_checkpoints_folder, "DUSt3R_ViTLarge_BaseDecoder_224_linear.pth")
+    extract_cross_attention_weights(
+        dust3r_ckpt_filepath, info_sharing_output_folder, "Two_View_Cross_Attention_Transformer_DUSt3R_224_linear.pth"
+    )
+    extract_dust3r_linear_checkpoints(dust3r_ckpt_filepath, pred_head_output_folder, "DUSt3R_224_linear")
+
+    # Extract dust3r 512 linear checkpoint
+    print("Extracting DUSt3R 512 linear checkpoint...")
+    dust3r_ckpt_filepath = os.path.join(args.dust3r_checkpoints_folder, "DUSt3R_ViTLarge_BaseDecoder_512_linear.pth")
+    extract_cross_attention_weights(
+        dust3r_ckpt_filepath, info_sharing_output_folder, "Two_View_Cross_Attention_Transformer_DUSt3R_512_linear.pth"
+    )
+    extract_dust3r_linear_checkpoints(dust3r_ckpt_filepath, pred_head_output_folder, "DUSt3R_512_linear")
+
+    # Extract dust3r 512 dpt checkpoint
+    print("Extracting DUSt3R 512 dpt checkpoint...")
+    dust3r_ckpt_filepath = os.path.join(args.dust3r_checkpoints_folder, "DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth")
+    extract_cross_attention_weights(
+        dust3r_ckpt_filepath, info_sharing_output_folder, "Two_View_Cross_Attention_Transformer_DUSt3R_512_dpt.pth"
+    )
+    extract_dust3r_dpt_checkpoints(dust3r_ckpt_filepath, pred_head_output_folder, "DUSt3R_512_dpt")
+
+    # Extract mast3r 512 dpt checkpoint
+    print("Extracting MASt3R 512 dpt checkpoint...")
+    mast3r_ckpt_path = os.path.join(
+        args.dust3r_checkpoints_folder, "MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth"
+    )
+    extract_cross_attention_weights(
+        mast3r_ckpt_path, info_sharing_output_folder, "Two_View_Cross_Attention_Transformer_MASt3R_512_dpt.pth"
+    )
+    extract_mast3r_dpt_checkpoints(mast3r_ckpt_path, pred_head_output_folder, "MASt3R_512_dpt")