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
140
141
142
143
144
145
146
147
148
149
150
|
from __future__ import annotations
import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
import json
from pathlib import Path
import time
from PIL import Image
import requests
import torch
from tqdm import tqdm
from transformers import AutoImageProcessor, AutoModel
from .common import dtype_for_device
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser()
p.add_argument(
"--nodes", default="artifacts/vg/vision_nodes.private.jsonl"
)
p.add_argument("--output", default="artifacts/vg/vision_features.pt")
p.add_argument(
"--image-cache", default="/tmp/yurenh2-worldalign-vg-images"
)
p.add_argument("--model", default="facebook/dinov2-small")
p.add_argument("--device", default="cuda:1")
p.add_argument("--node-batch-size", type=int, default=8)
p.add_argument("--download-workers", type=int, default=32)
p.add_argument("--limit", type=int)
return p.parse_args()
def read_jsonl(path: str) -> list[dict]:
with open(path, encoding="utf-8") as handle:
return [json.loads(line) for line in handle if line.strip()]
def download_one(record: dict, image_cache: Path) -> Path:
image_cache.mkdir(parents=True, exist_ok=True)
output = image_cache / f"{int(record['source_image_id'])}.jpg"
if output.exists() and output.stat().st_size > 0:
return output
temporary = output.with_suffix(".jpg.part")
last_error: Exception | None = None
for attempt in range(3):
try:
response = requests.get(record["url"], timeout=30)
response.raise_for_status()
temporary.write_bytes(response.content)
with Image.open(temporary) as image:
image.verify()
temporary.replace(output)
return output
except Exception as error:
last_error = error
if temporary.exists():
temporary.unlink()
time.sleep(1 + attempt)
raise RuntimeError(f"Failed to download {record['url']}") from last_error
def crop_views(record: dict, path: Path) -> tuple[Image.Image, list[Image.Image]]:
with Image.open(path) as source:
image = source.convert("RGB")
width, height = image.size
crops: list[Image.Image] = []
for region in record["regions"]:
x0 = max(0, min(int(region["x"]), width - 1))
y0 = max(0, min(int(region["y"]), height - 1))
x1 = max(x0 + 1, min(x0 + int(region["width"]), width))
y1 = max(y0 + 1, min(y0 + int(region["height"]), height))
crops.append(image.crop((x0, y0, x1, y1)))
return image, crops
@torch.inference_mode()
def main() -> None:
args = parse_args()
records = read_jsonl(args.nodes)
if args.limit:
records = records[: args.limit]
if not records:
raise ValueError("No vision nodes found")
view_count = len(records[0]["regions"])
if any(len(record["regions"]) != view_count for record in records):
raise ValueError("All nodes must have the same number of region views")
image_cache = Path(args.image_cache)
paths: dict[str, Path] = {}
with ThreadPoolExecutor(max_workers=args.download_workers) as executor:
futures = {
executor.submit(download_one, record, image_cache): record["node_id"]
for record in records
}
for future in tqdm(
as_completed(futures), total=len(futures), desc="VG image download"
):
paths[futures[future]] = future.result()
processor = AutoImageProcessor.from_pretrained(args.model)
dtype = dtype_for_device(args.device)
model = AutoModel.from_pretrained(args.model, torch_dtype=dtype).to(
args.device
)
model.eval()
region_outputs: list[torch.Tensor] = []
global_outputs: list[torch.Tensor] = []
for start in tqdm(
range(0, len(records), args.node_batch_size),
desc="DINO VG bundles",
):
batch = records[start : start + args.node_batch_size]
images: list[Image.Image] = []
for record in batch:
global_image, crops = crop_views(record, paths[record["node_id"]])
images.append(global_image)
images.extend(crops)
pixels = processor(images=images, return_tensors="pt")[
"pixel_values"
].to(args.device, dtype=dtype)
hidden = model(pixel_values=pixels, return_dict=True).last_hidden_state[
:, 0
]
hidden = hidden.float().cpu().reshape(
len(batch), view_count + 1, -1
)
global_outputs.append(hidden[:, 0])
region_outputs.append(hidden[:, 1:])
state = {
"model": args.model,
"node_ids": [record["node_id"] for record in records],
"region_features": torch.cat(region_outputs),
"global_features": torch.cat(global_outputs),
"views_per_node": view_count,
"source_metadata_removed": True,
}
Path(args.output).parent.mkdir(parents=True, exist_ok=True)
torch.save(state, args.output)
print(
f"Wrote {args.output}: regions={tuple(state['region_features'].shape)}, "
f"global={tuple(state['global_features'].shape)}"
)
if __name__ == "__main__":
main()
|
