File size: 9,215 Bytes
7dfc72e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
#!/usr/bin/env python3
"""Minimal GPU lane eval using offline precomputed tokens.

WARNING: This is a lightweight DEBUGGING HELPER only.
It is NOT equivalent to the primary online evaluation (eval_atlas.py).
Differences from main eval: uses offline tokens (no live encoders),
may fall back to legacy Chamfer if openlanev2 is missing, does not
auto-detect LoRA, may use different sampler defaults.
Do NOT use results from this script as official metrics.
For production evaluation use: bash scripts/eval_checkpoint.sh (online mode).
"""
import sys, os

if os.environ.get("ATLAS_ALLOW_OFFLINE", "").lower() not in ("1", "true", "yes"):
    print(
        "ERROR: This is an OFFLINE debugging helper, not the primary online evaluation.\n"
        "It is isolated by default to prevent accidental use in experiments.\n"
        "If you really need it, set:  ATLAS_ALLOW_OFFLINE=1\n"
        "For production evaluation use:  bash scripts/eval_checkpoint.sh",
        file=sys.stderr,
    )
    sys.exit(1)

import json, torch, numpy as np
from collections import defaultdict
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "3")

from src.model.modeling_atlas import AtlasForCausalLM
from src.model.topomlp_adapter import TopoMLPToAtlasMapTokens
from src.dataset.atlas_dataset import AtlasDataset, make_atlas_collate_fn, load_tokenizer
from src.dataset.scene_sampler import SceneSequentialSampler
from src.eval.metrics import parse_atlas_output
from train_atlas import _reconstruct_topomlp_outs

CKPT = sys.argv[1]
MAX_SAMPLES = int(sys.argv[2]) if len(sys.argv) > 2 else 20
OUT_JSON = sys.argv[3] if len(sys.argv) > 3 else None

DATA_JSON = "data/openlane_subsetB_lane_val_4pt.json"
DATA_ROOT = "/home/guoyuanbo/autodl-tmp/data/nuscenes"
DET_TOKENS = "work_dirs/precomputed_det_tokens_offline/val"
MAP_TOKENS = "work_dirs/precomputed_map_tokens_offline/val"
LLM = "pretrained/vicuna-7b-v1.5"

print(f"Checkpoint: {CKPT}", flush=True)
print(f"Max samples: {MAX_SAMPLES}", flush=True)
print(f"Task: lane detection", flush=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = load_tokenizer(LLM)
if "<query>" not in tokenizer.get_vocab():
    tokenizer.add_tokens(["<query>"])

atlas = AtlasForCausalLM(
    llm_model_name=LLM, visual_hidden_size=256,
    num_queries=256, num_map_queries=256,
    use_flash_attention=False, torch_dtype=torch.bfloat16,
)
atlas.resize_token_embeddings(len(tokenizer))
atlas.set_query_token_id(tokenizer.convert_tokens_to_ids("<query>"))

ckpt_data = torch.load(CKPT, map_location="cpu")
atlas.load_state_dict(ckpt_data["atlas_state_dict"], strict=False)

topomlp_adapter = TopoMLPToAtlasMapTokens(
    num_map_tokens=256, hidden_size=256,
    bev_range=(-51.2, -25.6, -8.0, 51.2, 25.6, 4.0),
).to(device)
if "adapter_state_dict" in ckpt_data and ckpt_data["adapter_state_dict"]:
    topomlp_adapter.load_state_dict(ckpt_data["adapter_state_dict"], strict=False)
    print("Loaded adapter weights from checkpoint", flush=True)
else:
    print("INFO: No adapter_state_dict in checkpoint. TopoMLPToAtlasMapTokens is a "
          "parameter-free Top-K selector; learned projection lives in "
          "atlas_state_dict (projector_map/rp). This is normal.", flush=True)
topomlp_adapter.eval()

del ckpt_data
atlas = atlas.to(device)
atlas.eval()
print("Model loaded on GPU (bf16)", flush=True)

dataset = AtlasDataset(
    json_file=DATA_JSON, image_root=DATA_ROOT, tokenizer=tokenizer,
    max_length=4096, is_training=False,
    precomputed_det_tokens=DET_TOKENS, precomputed_map_tokens=MAP_TOKENS,
)
scene_groups = dataset.get_scene_groups()
sampler = SceneSequentialSampler(scene_groups)
collate_fn = make_atlas_collate_fn(tokenizer.pad_token_id)
loader = torch.utils.data.DataLoader(
    dataset, batch_size=1, shuffle=False, sampler=sampler, num_workers=0, collate_fn=collate_fn,
)
print("Sampling: scene-sequential (paper-aligned)", flush=True)

try:
    from openlanev2.evaluation.f_score import LaneEval
    evaluator = LaneEval()
    USE_OFFICIAL = True
    print("Using OpenLane-V2 official F-Score evaluator", flush=True)
except ImportError:
    USE_OFFICIAL = False
    print("openlanev2 not available, using legacy Chamfer eval", flush=True)

all_pred_lanes = []
all_gt_lanes = []
count = 0
data_by_id = defaultdict(list)
for idx_i, item in enumerate(dataset.data):
    data_by_id[str(item.get("id", ""))].append((idx_i, item))

_params = list(topomlp_adapter.parameters())
adapter_dtype = _params[0].dtype if _params else torch.float32

for batch in loader:
    if count >= MAX_SAMPLES:
        break
    input_ids = batch["input_ids"].to(device)
    attention_mask = batch["attention_mask"].to(device)

    vis = {}
    if "precomputed_det" not in batch or "precomputed_det_ref" not in batch:
        raise RuntimeError(
            f"Precomputed det tokens missing for sample {batch.get('sample_id', ['?'])[0]}. "
            f"This offline helper requires precomputed tokens in {DET_TOKENS}."
        )
    vis["detection"] = batch["precomputed_det"].to(device)
    vis["detection_ref_points"] = batch["precomputed_det_ref"].to(device)

    if "precomputed_map" not in batch:
        raise RuntimeError(
            f"Precomputed map tokens missing for sample {batch.get('sample_id', ['?'])[0]}. "
            f"This offline helper requires precomputed tokens in {MAP_TOKENS}."
        )
    outs = _reconstruct_topomlp_outs(batch["precomputed_map"][0], device, adapter_dtype)
    with torch.no_grad():
        map_out = topomlp_adapter(outs)
    vis["map"] = map_out["map"]
    vis["map_ref_points"] = map_out["map_ref_points"]

    with torch.no_grad():
        gen = atlas.generate(
            input_ids=input_ids, attention_mask=attention_mask,
            visual_features=vis, max_new_tokens=2700, do_sample=False,
        )
    text_full = tokenizer.decode(gen[0], skip_special_tokens=True)
    text = text_full.split("ASSISTANT:")[-1].strip() if "ASSISTANT:" in text_full else text_full.strip()

    sample_id = str(batch["sample_id"][0]) if "sample_id" in batch else ""
    candidates = data_by_id.get(sample_id, [])
    if len(candidates) == 1:
        item = candidates[0][1]
    elif len(candidates) > 1:
        item = candidates[0][1]
    else:
        item = dataset.data[count]
    conv = item.get("conversations", [])
    gt_answer = ""
    for turn in conv:
        if turn.get("from") in ("gpt", "assistant"):
            gt_answer = turn.get("value", "")
            break

    preds = [p for p in parse_atlas_output(text) if p.get("type") == "lane"]
    gt_lanes = [p for p in parse_atlas_output(gt_answer) if p.get("type") == "lane"]
    all_pred_lanes.append(preds)
    all_gt_lanes.append(gt_lanes)

    count += 1
    if count % 5 == 0:
        print(f"  [{count}/{MAX_SAMPLES}] sample_id={sample_id} pred_lanes={len(preds)} gt_lanes={len(gt_lanes)}", flush=True)

if USE_OFFICIAL:
    def _to_ndarray_list(lanes):
        out = []
        for lane in lanes:
            pts = lane.get("points", [])
            if not pts:
                continue
            rows = []
            for pt in pts:
                if isinstance(pt, dict):
                    rows.append(pt.get("world_coords", [0, 0, 0])[:3])
                else:
                    rows.append(list(pt)[:3])
            arr = np.array(rows, dtype=np.float64)
            if arr.shape[0] >= 2:
                out.append(arr)
        return out

    stats = []
    for pl, gl in zip(all_pred_lanes, all_gt_lanes):
        pa = _to_ndarray_list(pl)
        ga = _to_ndarray_list(gl)
        pc = [np.int8(1)] * len(pa)
        gc = [np.int8(1)] * len(ga)
        r, p, c, ng, np_, mn = evaluator.bench(pa, pc, ga, gc)
        stats.append(np.array([r, p, c, ng, np_, mn]))
    if stats:
        s = np.array(stats)
        tg = np.sum(s[:, 3])
        tp = np.sum(s[:, 4])
        recall = np.sum(s[:, 0]) / max(tg, 1e-6)
        precision = np.sum(s[:, 1]) / max(tp, 1e-6)
        f1 = 2 * recall * precision / (recall + precision) if (recall + precision) > 0 else 0.0
    else:
        f1 = recall = precision = 0.0
else:
    from src.eval.metrics import calculate_lane_detection_metrics
    ttp = tfp = tfn = 0
    for pl, gl in zip(all_pred_lanes, all_gt_lanes):
        m = calculate_lane_detection_metrics(pl, gl)
        ttp += m["lane_tp"]; tfp += m["lane_fp"]; tfn += m["lane_fn"]
    precision = ttp / (ttp + tfp) if (ttp + tfp) > 0 else 0
    recall = ttp / (ttp + tfn) if (ttp + tfn) > 0 else 0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0

results = {
    "lane_f1": round(float(f1), 4),
    "lane_precision": round(float(precision), 4),
    "lane_recall": round(float(recall), 4),
    "num_samples": count,
}
print(f"\n=== Lane Results ===", flush=True)
for k, v in sorted(results.items()):
    print(f"  {k}: {v}", flush=True)

if OUT_JSON:
    with open(OUT_JSON, "w") as f:
        json.dump({"metrics": {"lane": results}, "num_samples": count,
                   "sampling": "scene_sequential",
                   "checkpoint": CKPT}, f, indent=2)
    print(f"Saved to {OUT_JSON}", flush=True)