236 lines
8.2 KiB
Python
236 lines
8.2 KiB
Python
"""
|
|
Deep-learning matchers: SuperPoint+LightGlue ("SuperGlue" in the UI/output
|
|
naming, per the existing convention -- LightGlue is the permissively
|
|
licensed, actively maintained successor actually running under the hood)
|
|
and LoFTR (dense, pairwise).
|
|
|
|
Template-side features/tensors are precomputed once at startup and reused
|
|
for every request; only the query image is processed per request. Since
|
|
there are only a handful of templates, even LoFTR (which has no reusable
|
|
per-image descriptor and must re-run a full forward pass per pair) is cheap
|
|
here -- it does NOT scale the way it would against hundreds of images.
|
|
"""
|
|
|
|
import gc
|
|
import logging
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
|
|
import config
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
_superpoint = None
|
|
_lightglue = None
|
|
_loftr = None
|
|
|
|
_template_superpoint_feats = {} # template_name -> {"feats", "mask"}
|
|
_template_loftr_tensors = {} # template_name -> {"tensor", "mask"}
|
|
|
|
|
|
def get_superpoint():
|
|
global _superpoint
|
|
if _superpoint is None:
|
|
logger.info("Loading SuperPoint on %s...", DEVICE)
|
|
from lightglue import SuperPoint
|
|
_superpoint = SuperPoint(
|
|
max_num_keypoints=config.SUPERPOINT_MAX_KEYPOINTS
|
|
).eval().to(DEVICE)
|
|
return _superpoint
|
|
|
|
|
|
def get_lightglue():
|
|
global _lightglue
|
|
if _lightglue is None:
|
|
logger.info("Loading LightGlue on %s...", DEVICE)
|
|
from lightglue import LightGlue
|
|
_lightglue = LightGlue(features="superpoint").eval().to(DEVICE)
|
|
return _lightglue
|
|
|
|
|
|
def get_loftr():
|
|
global _loftr
|
|
if _loftr is None:
|
|
logger.info("Loading LoFTR on %s...", DEVICE)
|
|
import kornia.feature as KF
|
|
_loftr = KF.LoFTR(pretrained="outdoor").eval().to(DEVICE)
|
|
return _loftr
|
|
|
|
|
|
def unload_models():
|
|
"""Drops the SuperPoint/LightGlue/LoFTR model objects (not the small
|
|
per-template feature/tensor caches, which stay put) and frees their CUDA
|
|
memory. They're plain lazy singletons (see get_superpoint/get_lightglue/
|
|
get_loftr above), so the next call to any of those simply reloads --
|
|
exactly like a fresh process start. Used to make room for SAM-based
|
|
flower counting on an 8GB card that can't hold everything at once."""
|
|
global _superpoint, _lightglue, _loftr
|
|
freed = _superpoint is not None or _lightglue is not None or _loftr is not None
|
|
_superpoint = None
|
|
_lightglue = None
|
|
_loftr = None
|
|
if freed:
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
return freed
|
|
|
|
|
|
def masked_gray_tensor(bgr, mask, max_dim=None, fill_value=0.5):
|
|
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY).astype(np.float32) / 255.0
|
|
|
|
if max_dim is not None:
|
|
h, w = gray.shape
|
|
scale = max_dim / max(h, w)
|
|
if scale < 1.0:
|
|
new_w, new_h = int(w * scale), int(h * scale)
|
|
gray = cv2.resize(gray, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
|
mask = cv2.resize(mask, (new_w, new_h), interpolation=cv2.INTER_NEAREST)
|
|
|
|
gray[mask == 0] = fill_value
|
|
tensor = torch.from_numpy(gray)[None, None].to(DEVICE)
|
|
return tensor, mask
|
|
|
|
|
|
def keep_points_inside_mask(pts, mask):
|
|
if len(pts) == 0:
|
|
return np.zeros(0, dtype=bool)
|
|
h, w = mask.shape
|
|
xs = np.clip(pts[:, 0].round().astype(int), 0, w - 1)
|
|
ys = np.clip(pts[:, 1].round().astype(int), 0, h - 1)
|
|
return mask[ys, xs] > 0
|
|
|
|
|
|
# ---------------------------------------------------------------
|
|
# SuperPoint + LightGlue
|
|
# ---------------------------------------------------------------
|
|
|
|
def superpoint_extract(bgr, mask):
|
|
# ASSUMPTION: cap at MAX_IMAGE_DIM even though uploads are already
|
|
# resized upstream -- this is what used to run at full (sometimes
|
|
# multi-thousand-pixel) upload resolution and was the main driver of
|
|
# the multi-GB memory spikes that got the process OOM-killed.
|
|
tensor, resized_mask = masked_gray_tensor(bgr, mask, max_dim=config.MAX_IMAGE_DIM)
|
|
with torch.no_grad():
|
|
feats = get_superpoint().extract(tensor)
|
|
return feats, resized_mask
|
|
|
|
|
|
def set_template_superpoint(template_name, bgr, mask):
|
|
feats, resized_mask = superpoint_extract(bgr, mask)
|
|
_template_superpoint_feats[template_name] = {"feats": feats, "mask": resized_mask}
|
|
|
|
|
|
def _lightglue_pair(feats_q, feats_t, mask_q, mask_t):
|
|
from lightglue.utils import rbd
|
|
|
|
with torch.no_grad():
|
|
matches01 = get_lightglue()({"image0": feats_q, "image1": feats_t})
|
|
|
|
feats_q_, feats_t_, matches01_ = [rbd(x) for x in [feats_q, feats_t, matches01]]
|
|
matches = matches01_["matches"]
|
|
if matches.shape[0] == 0:
|
|
return 0, 0.0
|
|
|
|
kpts_q = feats_q_["keypoints"][matches[..., 0]].cpu().numpy()
|
|
kpts_t = feats_t_["keypoints"][matches[..., 1]].cpu().numpy()
|
|
|
|
keep = keep_points_inside_mask(kpts_q, mask_q) & keep_points_inside_mask(kpts_t, mask_t)
|
|
kpts_q, kpts_t = kpts_q[keep], kpts_t[keep]
|
|
|
|
if len(kpts_q) < 4:
|
|
return 0, 0.0
|
|
|
|
src = kpts_q.reshape(-1, 1, 2).astype(np.float32)
|
|
dst = kpts_t.reshape(-1, 1, 2).astype(np.float32)
|
|
|
|
_, ransac_mask = cv2.findHomography(src, dst, cv2.RANSAC, 5.0)
|
|
if ransac_mask is None:
|
|
return 0, 0.0
|
|
|
|
inlier_count = int(ransac_mask.sum())
|
|
confidence_pct = (inlier_count / len(src)) * 100 if len(src) > 0 else 0.0
|
|
return inlier_count, confidence_pct
|
|
|
|
|
|
def superglue_match_against_templates(bgr, mask):
|
|
feats_q, mask_q = superpoint_extract(bgr, mask)
|
|
results = []
|
|
for template_name, data in _template_superpoint_feats.items():
|
|
inlier_count, confidence_pct = _lightglue_pair(
|
|
feats_q, data["feats"], mask_q, data["mask"]
|
|
)
|
|
results.append({
|
|
"template": template_name,
|
|
"score": inlier_count,
|
|
"confidence": round(confidence_pct, 1),
|
|
})
|
|
results.sort(key=lambda r: r["score"], reverse=True)
|
|
return results
|
|
|
|
|
|
# ---------------------------------------------------------------
|
|
# LoFTR
|
|
# ---------------------------------------------------------------
|
|
|
|
def set_template_loftr(template_name, bgr, mask):
|
|
tensor, resized_mask = masked_gray_tensor(bgr, mask, max_dim=config.LOFTR_MAX_DIM)
|
|
_template_loftr_tensors[template_name] = {"tensor": tensor, "mask": resized_mask}
|
|
|
|
|
|
def _loftr_pair(tensor_q, tensor_t, mask_q, mask_t):
|
|
with torch.no_grad():
|
|
out = get_loftr()({"image0": tensor_q, "image1": tensor_t})
|
|
|
|
conf = out["confidence"].cpu().numpy()
|
|
keep_conf = conf >= config.LOFTR_CONFIDENCE_THRESHOLD
|
|
|
|
kpts_q = out["keypoints0"].cpu().numpy()[keep_conf]
|
|
kpts_t = out["keypoints1"].cpu().numpy()[keep_conf]
|
|
|
|
keep = keep_points_inside_mask(kpts_q, mask_q) & keep_points_inside_mask(kpts_t, mask_t)
|
|
kpts_q, kpts_t = kpts_q[keep], kpts_t[keep]
|
|
|
|
if len(kpts_q) < 4:
|
|
return 0, 0.0
|
|
|
|
src = kpts_q.reshape(-1, 1, 2).astype(np.float32)
|
|
dst = kpts_t.reshape(-1, 1, 2).astype(np.float32)
|
|
|
|
_, ransac_mask = cv2.findHomography(src, dst, cv2.RANSAC, 5.0)
|
|
if ransac_mask is None:
|
|
return 0, 0.0
|
|
|
|
inlier_count = int(ransac_mask.sum())
|
|
confidence_pct = (inlier_count / len(src)) * 100 if len(src) > 0 else 0.0
|
|
return inlier_count, confidence_pct
|
|
|
|
|
|
def loftr_match_against_templates(bgr, mask):
|
|
tensor_q, mask_q = masked_gray_tensor(bgr, mask, max_dim=config.LOFTR_MAX_DIM)
|
|
results = []
|
|
for template_name, data in _template_loftr_tensors.items():
|
|
inlier_count, confidence_pct = _loftr_pair(
|
|
tensor_q, data["tensor"], mask_q, data["mask"]
|
|
)
|
|
results.append({
|
|
"template": template_name,
|
|
"score": inlier_count,
|
|
"confidence": round(confidence_pct, 1),
|
|
})
|
|
# LoFTR has no reusable per-image descriptor -- this is a full dense
|
|
# CNN+transformer forward pass PER TEMPLATE (6 of them per query,
|
|
# back to back). Left uncleared, each pass's intermediate
|
|
# activations fragment the allocator further than SuperPoint/
|
|
# LightGlue's much lighter per-template cost ever does, and on an
|
|
# 8GB card that's enough on its own to tip into CUDA OOM a few
|
|
# templates in.
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
results.sort(key=lambda r: r["score"], reverse=True)
|
|
return results
|