third_party/TopicFM/viz/methods/topicfm.py

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 VizTopicFM(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
        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 = "TopicFM"

        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))

        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
                )  # 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()