File size: 3,449 Bytes
989d4ef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | """Text -> image search with a trained Native-Bird x SigLIP alignment checkpoint.
python search.py --ckpt runs/native_bird/best.pt "Nkịta na-agba ọsọ n'ahịhịa" "A dog running on grass" -k 5
python search.py --ckpt runs/native_bird/best.pt --split test --json out.json "query 1" "query 2"
The gallery is every image of --split (default test, 809 images). Image features come from the same frozen
image encoder the checkpoint was trained with (cached in data/features/, computed on first use; needs the
Flickr8k images at data.image_dir). Works on CPU or GPU.
"""
import argparse, json, os, torch
from nbmm.common import (seed_all, pick_device, load_split, image_features, text_features, load_sentence_model,
text_arm, TrainableText, AlignModel, rel, _set)
class Searcher:
def __init__(self, ckpt, split="test", device="auto", overrides=()):
ck = torch.load(rel(ckpt), map_location="cpu", weights_only=False)
self.cfg = cfg = ck["cfg"]
for o in overrides:
k, v = o.split("=", 1); _set(cfg, k, v)
seed_all(cfg["seed"])
self.dev = pick_device(device)
self.df = load_split(cfg, split)
self.ids = self.df.image_id.tolist()
I = image_features(cfg, self.df, split, self.dev)
self.model = AlignModel(ck["d_img"], ck["d_txt"], cfg).to(self.dev).eval()
self.model.load_state_dict(ck["model"])
self.unfreeze = cfg["text_encoder"]["unfreeze_last_n_layers"] > 0
if self.unfreeze:
self.txt = TrainableText(cfg, self.dev).to(self.dev).eval()
self.txt.load_state_dict(ck["text_trainable"], strict=False)
else:
self.arm = text_arm(cfg)
self.st = load_sentence_model(self.arm, self.dev)
with torch.no_grad():
self.zi = self.model.img(I.to(self.dev))
@torch.no_grad()
def embed_text(self, texts):
if self.unfreeze:
e = self.txt(texts)
else:
e = self.st.encode([self.arm.get("prompt", "") + t for t in texts], convert_to_tensor=True,
normalize_embeddings=True).float().to(self.dev)
return self.model.txt(e)
@torch.no_grad()
def search(self, texts, k=5):
sims = self.embed_text(texts) @ self.zi.t()
out = []
for t, row in zip(texts, sims):
order = row.argsort(descending=True)
top = order[:k].tolist()
out.append({"query": t, "top": [{"image_id": self.ids[j], "score": round(row[j].item(), 4)} for j in top],
"_order": order})
return out
def rank_of(self, res, image_id):
j = self.ids.index(image_id)
return int((res["_order"] == j).nonzero()[0].item()) + 1 # 1-based
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--ckpt", required=True)
ap.add_argument("--split", default="test")
ap.add_argument("-k", type=int, default=5)
ap.add_argument("--json", default=None)
ap.add_argument("queries", nargs="+")
a = ap.parse_args()
s = Searcher(a.ckpt, a.split)
res = s.search(a.queries, a.k)
for r in res:
r.pop("_order")
print(r["query"])
for i, h in enumerate(r["top"], 1):
print(f" {i}. {h['image_id']} {h['score']}")
if a.json:
json.dump(res, open(a.json, "w"), indent=2, ensure_ascii=False)
if __name__ == "__main__":
main()
|