litert-community/xfeat-litert

XFeat (Accelerated Features) — LiteRT (CompiledModel GPU)

XFeat (Apache-2.0, ~1.5M, a lightweight pure-CNN local feature extractor for image matching — SLAM / AR / image registration) re-authored to a GPU-native LiteRT .tflite via litert_torch. FP16, 1.4 MB, input [1, 480, 640, 1] NHWC normalized grayscale.

167 mutual-nearest-neighbor matches between two views of the same scene, from the on-device fp16 model. Photo: "Lily the Golden Retriever in the grass" (Wikimedia Commons, Public Domain); second view is a synthetic homography (rotation + translation).

Verified on a Pixel 8a: full LITERT_CL residency (72/72 nodes, 1 partition), ~0.4 ms, GPU output matches CPU/PyTorch (corr 0.9999).

I/O

• Input [1, 480, 640, 1] NHWC, grayscale, per-image InstanceNorm applied host-side ((g - mean)/sqrt(var+1e-5) over the image).
• Outputs (all at H/8 × W/8 = 60×80): feats [1,64,60,80] dense descriptors; keypoints [1,65,60,80] keypoint logits; heatmap [1,1,60,80] reliability. Keypoint NMS, descriptor bilinear-sampling, and mutual-nearest-neighbor matching run host-side.

Minimal usage

import numpy as np
from PIL import Image
from ai_edge_litert.interpreter import Interpreter

def extract(path):
    g = np.asarray(Image.open(path).convert("L").resize((640, 480)), np.float32)
    g = (g - g.mean()) / np.sqrt(g.var() + 1e-5)          # host instance-norm
    it = Interpreter(model_path="xfeat_fp16.tflite"); it.allocate_tensors()
    it.set_tensor(it.get_input_details()[0]["index"], g[None, None]); it.invoke()
    feats, heat, klog = (it.get_tensor(o["index"])[0] for o in
                         sorted(it.get_output_details(), key=lambda o: o["index"]))
    return feats, heat, klog   # [64,60,80], [1,60,80], [65,60,80]
# decode: per-cell softmax over the 65 logits (64 positions + dustbin) * reliability,
# 5x5 NMS + top-K, bilinear-sample feats at kp/8, L2-normalize, mutual-NN (cos >= 0.82)

Kotlin (Android, LiteRT CompiledModel GPU)

// implementation("com.google.ai.edge.litert:litert:2.1.5")
val model = CompiledModel.create(File(ctx.filesDir, "xfeat_fp16.tflite").absolutePath,
    CompiledModel.Options(Accelerator.GPU), null)
val inBuf = model.createInputBuffers(); val outBuf = model.createOutputBuffers()
inBuf[0].writeFloat(grayNorm)             // [1,1,480,640] instance-normalized grayscale
model.run(inBuf, outBuf)
val feats = outBuf[0].readFloat()         // [64*60*80] dense descriptors
val heat = outBuf[1].readFloat()          // [60*80] reliability
val klog = outBuf[2].readFloat()          // [65*60*80] cell logits
// decode + mutual-NN matching: see XFeatMatcher.kt in the image_matching LiteRT sample.

GPU-clean re-authoring

• Input gray + InstanceNorm moved host-side (its spatial reduction over H·W would overflow fp16 on the delegate).
• _unfold2d(x, 8) (space-to-depth via unfold → >4-D / GATHER_ND) → a one-hot Conv2d(1,64,k=8,s=8) (exact, single CONV_2D). Result: zero GATHER/SELECT/TopK/Cast, no >4-D — full GPU residency.

Training data & PII

XFeat is trained on public correspondence data (MegaDepth + synthetic homographies). It outputs geometric keypoints/descriptors only — no faces, identities, or personal attributes. Official weights; only the op graph was re-authored for GPU.

Sample app + conversion script

https://github.com/google-ai-edge/litert-samples (compiled_model_api, two-image matching).