third_party/GlueStick/gluestick/models/wireframe.py

import numpy as np
import torch
from pytlsd import lsd
from sklearn.cluster import DBSCAN

from .base_model import BaseModel
from .superpoint import SuperPoint, sample_descriptors
from ..geometry import warp_lines_torch


def lines_to_wireframe(lines, line_scores, all_descs, conf):
    """Given a set of lines, their score and dense descriptors,
        merge close-by endpoints and compute a wireframe defined by
        its junctions and connectivity.
    Returns:
        junctions: list of [num_junc, 2] tensors listing all wireframe junctions
        junc_scores: list of [num_junc] tensors with the junction score
        junc_descs: list of [dim, num_junc] tensors with the junction descriptors
        connectivity: list of [num_junc, num_junc] bool arrays with True when 2 junctions are connected
        new_lines: the new set of [b_size, num_lines, 2, 2] lines
        lines_junc_idx: a [b_size, num_lines, 2] tensor with the indices of the junctions of each endpoint
        num_true_junctions: a list of the number of valid junctions for each image in the batch,
                            i.e. before filling with random ones
    """
    b_size, _, _, _ = all_descs.shape
    device = lines.device
    endpoints = lines.reshape(b_size, -1, 2)

    (
        junctions,
        junc_scores,
        junc_descs,
        connectivity,
        new_lines,
        lines_junc_idx,
        num_true_junctions,
    ) = ([], [], [], [], [], [], [])
    for bs in range(b_size):
        # Cluster the junctions that are close-by
        db = DBSCAN(eps=conf.nms_radius, min_samples=1).fit(endpoints[bs].cpu().numpy())
        clusters = db.labels_
        n_clusters = len(set(clusters))
        num_true_junctions.append(n_clusters)

        # Compute the average junction and score for each cluster
        clusters = torch.tensor(clusters, dtype=torch.long, device=device)
        new_junc = torch.zeros(n_clusters, 2, dtype=torch.float, device=device)
        new_junc.scatter_reduce_(
            0,
            clusters[:, None].repeat(1, 2),
            endpoints[bs],
            reduce="mean",
            include_self=False,
        )
        junctions.append(new_junc)
        new_scores = torch.zeros(n_clusters, dtype=torch.float, device=device)
        new_scores.scatter_reduce_(
            0,
            clusters,
            torch.repeat_interleave(line_scores[bs], 2),
            reduce="mean",
            include_self=False,
        )
        junc_scores.append(new_scores)

        # Compute the new lines
        new_lines.append(junctions[-1][clusters].reshape(-1, 2, 2))
        lines_junc_idx.append(clusters.reshape(-1, 2))

        # Compute the junction connectivity
        junc_connect = torch.eye(n_clusters, dtype=torch.bool, device=device)
        pairs = clusters.reshape(-1, 2)  # these pairs are connected by a line
        junc_connect[pairs[:, 0], pairs[:, 1]] = True
        junc_connect[pairs[:, 1], pairs[:, 0]] = True
        connectivity.append(junc_connect)

        # Interpolate the new junction descriptors
        junc_descs.append(
            sample_descriptors(junctions[-1][None], all_descs[bs : (bs + 1)], 8)[0]
        )

    new_lines = torch.stack(new_lines, dim=0)
    lines_junc_idx = torch.stack(lines_junc_idx, dim=0)
    return (
        junctions,
        junc_scores,
        junc_descs,
        connectivity,
        new_lines,
        lines_junc_idx,
        num_true_junctions,
    )


class SPWireframeDescriptor(BaseModel):
    default_conf = {
        "sp_params": {
            "has_detector": True,
            "has_descriptor": True,
            "descriptor_dim": 256,
            "trainable": False,
            # Inference
            "return_all": True,
            "sparse_outputs": True,
            "nms_radius": 4,
            "detection_threshold": 0.005,
            "max_num_keypoints": 1000,
            "force_num_keypoints": True,
            "remove_borders": 4,
        },
        "wireframe_params": {
            "merge_points": True,
            "merge_line_endpoints": True,
            "nms_radius": 3,
            "max_n_junctions": 500,
        },
        "max_n_lines": 250,
        "min_length": 15,
    }
    required_data_keys = ["image"]

    def _init(self, conf):
        self.conf = conf
        self.sp = SuperPoint(conf.sp_params)

    def detect_lsd_lines(self, x, max_n_lines=None):
        if max_n_lines is None:
            max_n_lines = self.conf.max_n_lines
        lines, scores, valid_lines = [], [], []
        for b in range(len(x)):
            # For each image on batch
            img = (x[b].squeeze().cpu().numpy() * 255).astype(np.uint8)
            if max_n_lines is None:
                b_segs = lsd(img)
            else:
                for s in [0.3, 0.4, 0.5, 0.7, 0.8, 1.0]:
                    b_segs = lsd(img, scale=s)
                    if len(b_segs) >= max_n_lines:
                        break

            segs_length = np.linalg.norm(b_segs[:, 2:4] - b_segs[:, 0:2], axis=1)
            # Remove short lines
            b_segs = b_segs[segs_length >= self.conf.min_length]
            segs_length = segs_length[segs_length >= self.conf.min_length]
            b_scores = b_segs[:, -1] * np.sqrt(segs_length)
            # Take the most relevant segments with
            indices = np.argsort(-b_scores)
            if max_n_lines is not None:
                indices = indices[:max_n_lines]
            lines.append(torch.from_numpy(b_segs[indices, :4].reshape(-1, 2, 2)))
            scores.append(torch.from_numpy(b_scores[indices]))
            valid_lines.append(torch.ones_like(scores[-1], dtype=torch.bool))

        lines = torch.stack(lines).to(x)
        scores = torch.stack(scores).to(x)
        valid_lines = torch.stack(valid_lines).to(x.device)
        return lines, scores, valid_lines

    def _forward(self, data):
        b_size, _, h, w = data["image"].shape
        device = data["image"].device

        if not self.conf.sp_params.force_num_keypoints:
            assert b_size == 1, "Only batch size of 1 accepted for non padded inputs"

        # Line detection
        if "lines" not in data or "line_scores" not in data:
            if "original_img" in data:
                # Detect more lines, because when projecting them to the image most of them will be discarded
                lines, line_scores, valid_lines = self.detect_lsd_lines(
                    data["original_img"], self.conf.max_n_lines * 3
                )
                # Apply the same transformation that is applied in homography_adaptation
                lines, valid_lines2 = warp_lines_torch(
                    lines, data["H"], False, data["image"].shape[-2:]
                )
                valid_lines = valid_lines & valid_lines2
                lines[~valid_lines] = -1
                line_scores[~valid_lines] = 0
                # Re-sort the line segments to pick the ones that are inside the image and have bigger score
                sorted_scores, sorting_indices = torch.sort(
                    line_scores, dim=-1, descending=True
                )
                line_scores = sorted_scores[:, : self.conf.max_n_lines]
                sorting_indices = sorting_indices[:, : self.conf.max_n_lines]
                lines = torch.take_along_dim(lines, sorting_indices[..., None, None], 1)
                valid_lines = torch.take_along_dim(valid_lines, sorting_indices, 1)
            else:
                lines, line_scores, valid_lines = self.detect_lsd_lines(data["image"])

        else:
            lines, line_scores, valid_lines = (
                data["lines"],
                data["line_scores"],
                data["valid_lines"],
            )
        if line_scores.shape[-1] != 0:
            line_scores /= (
                line_scores.new_tensor(1e-8) + line_scores.max(dim=1).values[:, None]
            )

        # SuperPoint prediction
        pred = self.sp(data)

        # Remove keypoints that are too close to line endpoints
        if self.conf.wireframe_params.merge_points:
            kp = pred["keypoints"]
            line_endpts = lines.reshape(b_size, -1, 2)
            dist_pt_lines = torch.norm(kp[:, :, None] - line_endpts[:, None], dim=-1)
            # For each keypoint, mark it as valid or to remove
            pts_to_remove = torch.any(
                dist_pt_lines < self.conf.sp_params.nms_radius, dim=2
            )
            # Simply remove them (we assume batch_size = 1 here)
            assert len(kp) == 1
            pred["keypoints"] = pred["keypoints"][0][~pts_to_remove[0]][None]
            pred["keypoint_scores"] = pred["keypoint_scores"][0][~pts_to_remove[0]][
                None
            ]
            pred["descriptors"] = pred["descriptors"][0].T[~pts_to_remove[0]].T[None]

        # Connect the lines together to form a wireframe
        orig_lines = lines.clone()
        if self.conf.wireframe_params.merge_line_endpoints and len(lines[0]) > 0:
            # Merge first close-by endpoints to connect lines
            (
                line_points,
                line_pts_scores,
                line_descs,
                line_association,
                lines,
                lines_junc_idx,
                num_true_junctions,
            ) = lines_to_wireframe(
                lines,
                line_scores,
                pred["all_descriptors"],
                conf=self.conf.wireframe_params,
            )

            # Add the keypoints to the junctions and fill the rest with random keypoints
            (all_points, all_scores, all_descs, pl_associativity) = [], [], [], []
            for bs in range(b_size):
                all_points.append(
                    torch.cat([line_points[bs], pred["keypoints"][bs]], dim=0)
                )
                all_scores.append(
                    torch.cat([line_pts_scores[bs], pred["keypoint_scores"][bs]], dim=0)
                )
                all_descs.append(
                    torch.cat([line_descs[bs], pred["descriptors"][bs]], dim=1)
                )

                associativity = torch.eye(
                    len(all_points[-1]), dtype=torch.bool, device=device
                )
                associativity[
                    : num_true_junctions[bs], : num_true_junctions[bs]
                ] = line_association[bs][
                    : num_true_junctions[bs], : num_true_junctions[bs]
                ]
                pl_associativity.append(associativity)

            all_points = torch.stack(all_points, dim=0)
            all_scores = torch.stack(all_scores, dim=0)
            all_descs = torch.stack(all_descs, dim=0)
            pl_associativity = torch.stack(pl_associativity, dim=0)
        else:
            # Lines are independent
            all_points = torch.cat(
                [lines.reshape(b_size, -1, 2), pred["keypoints"]], dim=1
            )
            n_pts = all_points.shape[1]
            num_lines = lines.shape[1]
            num_true_junctions = [num_lines * 2] * b_size
            all_scores = torch.cat(
                [
                    torch.repeat_interleave(line_scores, 2, dim=1),
                    pred["keypoint_scores"],
                ],
                dim=1,
            )
            pred["line_descriptors"] = self.endpoints_pooling(
                lines, pred["all_descriptors"], (h, w)
            )
            all_descs = torch.cat(
                [
                    pred["line_descriptors"].reshape(
                        b_size, self.conf.sp_params.descriptor_dim, -1
                    ),
                    pred["descriptors"],
                ],
                dim=2,
            )
            pl_associativity = torch.eye(n_pts, dtype=torch.bool, device=device)[
                None
            ].repeat(b_size, 1, 1)
            lines_junc_idx = (
                torch.arange(num_lines * 2, device=device)
                .reshape(1, -1, 2)
                .repeat(b_size, 1, 1)
            )

        del pred["all_descriptors"]  # Remove dense descriptors to save memory
        torch.cuda.empty_cache()

        return {
            "keypoints": all_points,
            "keypoint_scores": all_scores,
            "descriptors": all_descs,
            "pl_associativity": pl_associativity,
            "num_junctions": torch.tensor(num_true_junctions),
            "lines": lines,
            "orig_lines": orig_lines,
            "lines_junc_idx": lines_junc_idx,
            "line_scores": line_scores,
            "valid_lines": valid_lines,
        }

    @staticmethod
    def endpoints_pooling(segs, all_descriptors, img_shape):
        assert segs.ndim == 4 and segs.shape[-2:] == (2, 2)
        filter_shape = all_descriptors.shape[-2:]
        scale_x = filter_shape[1] / img_shape[1]
        scale_y = filter_shape[0] / img_shape[0]

        scaled_segs = torch.round(
            segs * torch.tensor([scale_x, scale_y]).to(segs)
        ).long()
        scaled_segs[..., 0] = torch.clip(scaled_segs[..., 0], 0, filter_shape[1] - 1)
        scaled_segs[..., 1] = torch.clip(scaled_segs[..., 1], 0, filter_shape[0] - 1)
        line_descriptors = [
            all_descriptors[
                None,
                b,
                ...,
                torch.squeeze(b_segs[..., 1]),
                torch.squeeze(b_segs[..., 0]),
            ]
            for b, b_segs in enumerate(scaled_segs)
        ]
        line_descriptors = torch.cat(line_descriptors)
        return line_descriptors  # Shape (1, 256, 308, 2)

    def loss(self, pred, data):
        raise NotImplementedError

    def metrics(self, pred, data):
        return {}