| """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 |
|
|
|
|
| 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() |
|
|
