Spaces:
Sleeping
Sleeping
File size: 4,613 Bytes
61ab00e 9ce0bb8 61ab00e 9ce0bb8 61ab00e deebb21 61ab00e deebb21 61ab00e deebb21 61ab00e | 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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | """Build a nearest-neighbor sentence index from 3D CNN clip embeddings."""
from __future__ import annotations
import argparse
from pathlib import Path
import sys
import numpy as np
PROJECT_ROOT = Path(__file__).resolve().parents[1]
SRC_ROOT = PROJECT_ROOT / "src"
if str(SRC_ROOT) not in sys.path:
sys.path.insert(0, str(SRC_ROOT))
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Build a 3D CNN sentence embedding index.")
parser.add_argument(
"--manifest",
default="data/processed/how2sign_sentences_top25.normalized.frames.jsonl",
help="Clip manifest JSONL with sampled_frames.",
)
parser.add_argument(
"--model",
default="models/cnn-3d-sentence-top25-normalized.keras",
help="Path to the trained Keras sentence CNN.",
)
parser.add_argument(
"--output",
default=None,
help="Where to save the .index.npz file. Defaults next to the model.",
)
parser.add_argument(
"--split",
default="all",
choices=("all", "train", "val", "test"),
help="Which manifest split to index. Use 'all' for the full support set.",
)
parser.add_argument(
"--top-k",
type=int,
default=3,
help="Voting neighborhood size for retrieval.",
)
parser.add_argument(
"--candidate-pool",
type=int,
default=10,
help="How many nearest clips to scan when ranking unique labels.",
)
parser.add_argument(
"--batch-size",
type=int,
default=16,
help="Batch size for embedding extraction.",
)
return parser
def main() -> None:
args = build_parser().parse_args()
from bridgelink_asl.clip_dataset import load_clip_dataset
from bridgelink_asl.cnn import select_frame_paths
from bridgelink_asl.sentence_inference import load_sentence_runtime
import tensorflow as tf
manifest_path = Path(args.manifest).expanduser().resolve()
model_path = Path(args.model).expanduser().resolve()
output_path = (
Path(args.output).expanduser().resolve()
if args.output
else model_path.with_suffix(".index.npz")
)
runtime = load_sentence_runtime(local_model_path=model_path)
if runtime.embedding_model is None:
raise RuntimeError("The sentence model does not expose a clip_embedding layer.")
records = load_clip_dataset(manifest_path)
if args.split != "all":
records = [record for record in records if record.split == args.split]
if not records:
raise RuntimeError(f"No records found for split={args.split!r} in {manifest_path}.")
clip_ids = [record.clip_id for record in records]
labels = [record.label for record in records]
selected_paths = [
[str(path) for path in select_frame_paths(record.sampled_frames, runtime.frame_count)]
for record in records
]
def _load_clip(paths):
images = tf.map_fn(
lambda path: _load_image(tf, path, runtime.image_size, runtime.channels),
paths,
fn_output_signature=tf.float32,
)
return images
dataset = (
tf.data.Dataset.from_tensor_slices(selected_paths)
.map(_load_clip, num_parallel_calls=tf.data.AUTOTUNE)
.batch(args.batch_size)
.prefetch(tf.data.AUTOTUNE)
)
embeddings = np.asarray(runtime.embedding_model.predict(dataset, verbose=0), dtype=np.float32)
norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
norms = np.where(norms == 0.0, 1.0, norms)
normalized = embeddings / norms
output_path.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(
output_path,
normalized_embeddings=normalized.astype(np.float32),
labels=np.asarray(labels, dtype="<U128"),
clip_ids=np.asarray(clip_ids, dtype="<U128"),
source_manifest=np.asarray(str(manifest_path), dtype="<U512"),
support_split=np.asarray(args.split, dtype="<U32"),
top_k=np.asarray(args.top_k, dtype=np.int32),
candidate_pool=np.asarray(args.candidate_pool, dtype=np.int32),
)
print(f"Wrote sentence embedding index to {output_path}")
print(f"Indexed clips: {len(records)}")
def _load_image(tf, path, image_size: int, channels: int):
image_bytes = tf.io.read_file(path)
image = tf.io.decode_image(image_bytes, channels=channels, expand_animations=False)
image = tf.image.resize(image, (image_size, image_size))
return tf.cast(image, tf.float32)
if __name__ == "__main__":
main()
|