Skip to content

Commit

Permalink
this black setup is crazy
Browse files Browse the repository at this point in the history
  • Loading branch information
ducha-aiki committed Jan 15, 2024
1 parent 37704d5 commit 51924d8
Showing 1 changed file with 3 additions and 4 deletions.
7 changes: 3 additions & 4 deletions lightglue/dog_hardnet.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,7 +188,7 @@ def extract_single_image(self, image: torch.Tensor):
pred["keypoint_scores"],
)
pred2 = {k: v[keep] for k, v in pred.items()}
if pred2['keypoints'] is not None:
if pred2["keypoints"] is not None:
pred = pred2
pred = {k: torch.from_numpy(v).float() for k, v in pred.items()}
if scores is not None:
Expand All @@ -200,8 +200,7 @@ def extract_single_image(self, image: torch.Tensor):
lafs = laf_from_center_scale_ori(
pred["keypoints"].reshape(1, -1, 2),
6.0 * pred["scales"].reshape(1, -1, 1, 1),
torch.rad2deg(pred["oris"]).reshape(1, -1, 1)
).to(device)
torch.rad2deg(pred["oris"]).reshape(1, -1, 1)).to(device)
self.laf_desc = self.laf_desc.to(device)
self.laf_desc.descriptor = self.laf_desc.descriptor.eval()
pred["descriptors"] = self.laf_desc(image[None], lafs).reshape(-1, 128)
Expand All @@ -218,7 +217,7 @@ def _forward(self, data: dict) -> dict:
img = image[k]
if im_size is not None:
w, h = data["image_size"][k]
img = img[:, :h.to(torch.int32), :w.to(torch.int32)]
img = img[:, : h.to(torch.int32), : w.to(torch.int32)]
p = self.extract_single_image(img)
pred.append(p)
pred = {k: torch.stack([p[k] for p in pred], 0).to(device) for k in pred[0]}
Expand Down

0 comments on commit 51924d8

Please sign in to comment.