File size: 2,443 Bytes
7002f4e | 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 | """Extract RemoteCLIP features and image-text similarity scores."""
import importlib.util
from pathlib import Path
import numpy as np
import torch
import yaml
ROOT = Path(__file__).resolve().parents[1]
def load_model_class():
spec = importlib.util.spec_from_file_location("remoteclip_model", ROOT / "model" / "remoteclip.py")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module.RemoteCLIP
def main():
with (ROOT / "conf" / "config.yaml").open(encoding="utf-8") as handle:
config = yaml.safe_load(handle)
checkpoint_path = ROOT / config["paths"]["checkpoint"]
data_path = ROOT / config["data"]["path"]
if not checkpoint_path.exists():
raise FileNotFoundError("Missing checkpoint. Run `python scripts/train.py` first.")
if not data_path.exists():
raise FileNotFoundError("Missing data. Run `python scripts/fake_data.py` first.")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
RemoteCLIP = load_model_class()
model = RemoteCLIP(
vocabulary_size=config["data"]["vocabulary_size"],
context_length=config["data"]["context_length"],
**config["model"],
).to(device)
model.load_state_dict(checkpoint["model"])
model.eval()
archive = np.load(data_path)
images = torch.from_numpy(archive["test_images"]).to(device)
tokens = torch.from_numpy(archive["test_tokens"]).to(device)
with torch.inference_mode():
image_features = model.encode_image(images)
text_features = model.encode_text(tokens)
similarities = image_features @ text_features.t()
output_dir = ROOT / config["paths"]["inference_dir"]
output_dir.mkdir(parents=True, exist_ok=True)
np.savez_compressed(
output_dir / "retrieval.npz",
similarities=similarities.cpu().numpy(),
image_features=image_features.cpu().numpy(),
text_features=text_features.cpu().numpy(),
labels=archive["test_labels"],
images=archive["test_images"],
data_source=archive["data_source"],
protocol=archive["protocol"],
)
print(
f"output={output_dir.relative_to(ROOT)} samples={len(images)} "
f"data_source={str(archive['data_source'])} protocol={str(archive['protocol'])}"
)
if __name__ == "__main__":
main()
|