Sasi_Vardhan_Vemuri/Vase-Matcher
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add project files

Browse Source
This commit is contained in:
Suman
2026-08-04 17:09:29 +05:30
parent 464a68cceb
commit ac76cc86bc
372 changed files with 12221 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

235
pipeline/deep.py Normal file
View File

@@ -0,0 +1,235 @@
"""
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