diff --git a/demo_superpoint.py b/demo_superpoint.py index 55e878f..f918de1 100755 --- a/demo_superpoint.py +++ b/demo_superpoint.py @@ -1,4 +1,4 @@ -#!/usr/bin/env python +#!/usr/bin/env python3 # # %BANNER_BEGIN% # --------------------------------------------------------------------- @@ -152,19 +152,19 @@ def nms_fast(self, in_corners, H, W, dist_thresh): """ Run a faster approximate Non-Max-Suppression on numpy corners shaped: 3xN [x_i,y_i,conf_i]^T - + Algo summary: Create a grid sized HxW. Assign each corner location a 1, rest are zeros. Iterate through all the 1's and convert them either to -1 or 0. Suppress points by setting nearby values to 0. - + Grid Value Legend: -1 : Kept. 0 : Empty or suppressed. 1 : To be processed (converted to either kept or supressed). - + NOTE: The NMS first rounds points to integers, so NMS distance might not be exactly dist_thresh. It also assumes points are within image boundaries. - + Inputs in_corners - 3xN numpy array with corners [x_i, y_i, confidence_i]^T. H - Image height. @@ -248,13 +248,13 @@ def run(self, img): heatmap = np.reshape(nodust, [Hc, Wc, self.cell, self.cell]) heatmap = np.transpose(heatmap, [0, 2, 1, 3]) heatmap = np.reshape(heatmap, [Hc*self.cell, Wc*self.cell]) - xs, ys = np.where(heatmap >= self.conf_thresh) # Confidence threshold. + ys, xs = np.where(heatmap >= self.conf_thresh) # Confidence threshold. if len(xs) == 0: return np.zeros((3, 0)), None, None pts = np.zeros((3, len(xs))) # Populate point data sized 3xN. - pts[0, :] = ys - pts[1, :] = xs - pts[2, :] = heatmap[xs, ys] + pts[0, :] = xs + pts[1, :] = ys + pts[2, :] = heatmap[ys, xs] pts, _ = self.nms_fast(pts, H, W, dist_thresh=self.nms_dist) # Apply NMS. inds = np.argsort(pts[2,:]) pts = pts[:,inds[::-1]] # Sort by confidence.