File size: 13,697 Bytes
a80d6bb
 
 
 
 
 
 
 
 
 
 
c74a070
a80d6bb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c74a070
 
 
 
 
 
 
 
 
a80d6bb
 
c74a070
a80d6bb
 
 
 
 
c74a070
 
 
 
 
 
 
 
 
a80d6bb
 
 
c74a070
 
 
 
 
 
a80d6bb
 
 
 
 
 
 
c74a070
a80d6bb
 
 
 
 
 
c74a070
 
 
a80d6bb
 
 
c74a070
 
 
 
 
 
 
 
 
a80d6bb
 
 
 
c74a070
 
 
 
 
a80d6bb
c74a070
 
 
 
 
 
 
a80d6bb
c74a070
 
 
 
 
a80d6bb
c74a070
 
a80d6bb
c74a070
a80d6bb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c74a070
 
a80d6bb
 
 
 
 
c74a070
 
a80d6bb
 
c74a070
 
a80d6bb
c74a070
 
 
a80d6bb
 
 
 
c74a070
 
 
 
 
a80d6bb
 
 
c74a070
a80d6bb
 
c74a070
 
 
 
 
a80d6bb
c74a070
 
 
a80d6bb
 
 
 
 
 
c74a070
a80d6bb
c74a070
a80d6bb
 
c74a070
 
a80d6bb
 
c74a070
 
 
 
 
a80d6bb
 
 
 
 
c74a070
 
 
 
 
 
 
 
 
 
 
 
 
 
a80d6bb
 
c74a070
a80d6bb
c74a070
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a80d6bb
 
 
 
 
 
 
 
c74a070
 
 
a80d6bb
 
 
c74a070
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a80d6bb
 
c74a070
 
 
 
 
 
 
 
 
 
 
 
a80d6bb
 
 
 
 
 
 
 
c74a070
 
 
a80d6bb
 
c74a070
 
 
 
 
 
 
 
 
 
a80d6bb
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
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 {}