Rewrite: GatedFusion architecture + UAV-GeoLoc dataset
Architecture v2: - Query branch: drone + text -> GatedFusion -> proj -> query_emb - Gallery branch: satellite -> proj -> gallery_emb - Single InfoNCE loss (asymmetric 0.6/0.4) - GatedFusion: learnable gated addition (sigma(alpha)*img + (1-sigma(alpha))*text) - Baseline mode: gate=1.0 (text ignored) Dataset: - UAV-GeoLoc loader with template captions from path metadata - 27 terrain types with predefined features - Random positive crop sampling per epoch Configs: balanced (gate=0.7), baseline (no text), text_heavy (gate=0.3) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
@@ -1,9 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""Evaluation utilities for caption quality test.
|
||||
"""Evaluation for caption quality test.
|
||||
|
||||
Implements retrieval metrics across four directions and a
|
||||
`delta_r_at_1` helper that compares caption-aware vs. image-only runs.
|
||||
Recall@K for query(drone+text) -> gallery(satellite).
|
||||
delta_r_at_1 compares caption-aware vs baseline runs.
|
||||
"""
|
||||
|
||||
import json
|
||||
@@ -23,19 +23,10 @@ def _recall_at_k(
|
||||
similarity: torch.Tensor,
|
||||
k_values: tuple[int, ...] = (1, 5, 10),
|
||||
) -> dict[int, float]:
|
||||
"""Compute Recall@K assuming positives on the diagonal.
|
||||
|
||||
Args:
|
||||
similarity: Pairwise similarity matrix [N_query, N_gallery].
|
||||
k_values: Tuple of K values to compute.
|
||||
|
||||
Returns:
|
||||
Dict mapping K -> recall in [0, 1].
|
||||
"""
|
||||
"""Recall@K assuming positives on the diagonal."""
|
||||
n_query = similarity.size(0)
|
||||
targets = torch.arange(n_query, device=similarity.device)
|
||||
sorted_idx = similarity.argsort(dim=1, descending=True)
|
||||
|
||||
result: dict[int, float] = {}
|
||||
for k in k_values:
|
||||
top_k = sorted_idx[:, :k]
|
||||
@@ -49,51 +40,29 @@ def _encode_dataset(
|
||||
model: DualEncoderCaptionTest,
|
||||
loader: DataLoader,
|
||||
device: str,
|
||||
include_captions: bool,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Encode every sample in the loader into the shared embedding space.
|
||||
|
||||
Args:
|
||||
model: Trained dual encoder.
|
||||
loader: DataLoader yielding collated batches.
|
||||
device: Target device string.
|
||||
include_captions: If False, caption embeddings are skipped.
|
||||
|
||||
Returns:
|
||||
Dict with keys 'drone', 'sat', 'cap_drone', 'cap_sat' -> [N, D].
|
||||
"""
|
||||
"""Encode all samples into query and gallery embeddings."""
|
||||
model.eval()
|
||||
all_drone: list[torch.Tensor] = []
|
||||
all_sat: list[torch.Tensor] = []
|
||||
all_cap_drone: list[torch.Tensor] = []
|
||||
all_cap_sat: list[torch.Tensor] = []
|
||||
all_query: list[torch.Tensor] = []
|
||||
all_gallery: list[torch.Tensor] = []
|
||||
|
||||
for batch in loader:
|
||||
drone_img = batch["drone_img"].to(device, non_blocking=True)
|
||||
sat_img = batch["sat_img"].to(device, non_blocking=True)
|
||||
captions_drone = batch["caption_drone"] if include_captions else None
|
||||
captions_sat = batch["caption_sat"] if include_captions else None
|
||||
caption_drone = batch["caption_drone"]
|
||||
|
||||
embeddings = model(
|
||||
drone_img=drone_img,
|
||||
sat_img=sat_img,
|
||||
caption_drone=captions_drone,
|
||||
caption_sat=captions_sat,
|
||||
caption_drone=caption_drone,
|
||||
)
|
||||
all_drone.append(embeddings["drone"].cpu())
|
||||
all_sat.append(embeddings["sat"].cpu())
|
||||
if include_captions:
|
||||
all_cap_drone.append(embeddings["cap_drone"].cpu())
|
||||
all_cap_sat.append(embeddings["cap_sat"].cpu())
|
||||
all_query.append(embeddings["query"].cpu())
|
||||
all_gallery.append(embeddings["gallery"].cpu())
|
||||
|
||||
out = {
|
||||
"drone": torch.cat(all_drone, dim=0),
|
||||
"sat": torch.cat(all_sat, dim=0),
|
||||
return {
|
||||
"query": torch.cat(all_query, dim=0),
|
||||
"gallery": torch.cat(all_gallery, dim=0),
|
||||
}
|
||||
if include_captions:
|
||||
out["cap_drone"] = torch.cat(all_cap_drone, dim=0)
|
||||
out["cap_sat"] = torch.cat(all_cap_sat, dim=0)
|
||||
return out
|
||||
|
||||
|
||||
def evaluate_retrieval(
|
||||
@@ -101,51 +70,25 @@ def evaluate_retrieval(
|
||||
loader: DataLoader,
|
||||
device: str,
|
||||
k_values: tuple[int, ...] = (1, 5, 10),
|
||||
include_captions: bool = True,
|
||||
) -> dict[str, float]:
|
||||
"""Compute retrieval metrics across four directions.
|
||||
|
||||
Directions reported (when captions included):
|
||||
drone -> sat, sat -> drone, text -> sat, text -> drone.
|
||||
|
||||
Args:
|
||||
model: Trained DualEncoderCaptionTest.
|
||||
loader: DataLoader over evaluation split.
|
||||
device: torch device string.
|
||||
k_values: Recall@K cutoffs.
|
||||
include_captions: If False, only image-image directions computed.
|
||||
"""Compute R@K for query->gallery and gallery->query.
|
||||
|
||||
Returns:
|
||||
Flat dict with keys like 'r@1_drone_to_sat', 'r@5_text_to_sat', etc.
|
||||
Flat dict: r@1_query_to_gallery, r@5_query_to_gallery, etc.
|
||||
"""
|
||||
feats = _encode_dataset(
|
||||
model=model,
|
||||
loader=loader,
|
||||
device=device,
|
||||
include_captions=include_captions,
|
||||
)
|
||||
feats = _encode_dataset(model=model, loader=loader, device=device)
|
||||
|
||||
metrics: dict[str, float] = {}
|
||||
|
||||
sim_d2s = feats["drone"] @ feats["sat"].t()
|
||||
sim_s2d = sim_d2s.t()
|
||||
sim_q2g = feats["query"] @ feats["gallery"].t()
|
||||
|
||||
for k, val in _recall_at_k(sim_d2s, k_values).items():
|
||||
metrics[f"r@{k}_drone_to_sat"] = val
|
||||
for k, val in _recall_at_k(sim_s2d, k_values).items():
|
||||
metrics[f"r@{k}_sat_to_drone"] = val
|
||||
for k, val in _recall_at_k(sim_q2g, k_values).items():
|
||||
metrics[f"r@{k}_query_to_gallery"] = val
|
||||
for k, val in _recall_at_k(sim_q2g.t(), k_values).items():
|
||||
metrics[f"r@{k}_gallery_to_query"] = val
|
||||
|
||||
if include_captions and "cap_sat" in feats and "cap_drone" in feats:
|
||||
sim_t2s = feats["cap_sat"] @ feats["sat"].t()
|
||||
sim_t2d = feats["cap_drone"] @ feats["drone"].t()
|
||||
sim_tcd2tcs = feats["cap_drone"] @ feats["cap_sat"].t()
|
||||
|
||||
for k, val in _recall_at_k(sim_t2s, k_values).items():
|
||||
metrics[f"r@{k}_text_to_sat"] = val
|
||||
for k, val in _recall_at_k(sim_t2d, k_values).items():
|
||||
metrics[f"r@{k}_text_to_drone"] = val
|
||||
for k, val in _recall_at_k(sim_tcd2tcs, k_values).items():
|
||||
metrics[f"r@{k}_capdrone_to_capsat"] = val
|
||||
# Gate value for diagnostics.
|
||||
metrics["gate"] = model.fusion.gate_value
|
||||
|
||||
return metrics
|
||||
|
||||
@@ -153,64 +96,36 @@ def evaluate_retrieval(
|
||||
def delta_r_at_1(
|
||||
full_metrics: dict[str, float],
|
||||
baseline_metrics: dict[str, float],
|
||||
direction: str = "drone_to_sat",
|
||||
) -> float:
|
||||
"""Compute caption-quality proxy: R@1 gain from adding captions.
|
||||
|
||||
Args:
|
||||
full_metrics: Metrics from training WITH caption losses.
|
||||
baseline_metrics: Metrics from training WITHOUT caption losses.
|
||||
direction: Retrieval direction to compare.
|
||||
|
||||
Returns:
|
||||
Δ R@1 in [−1, +1] range (positive = captions help).
|
||||
"""
|
||||
key = f"r@1_{direction}"
|
||||
if key not in full_metrics or key not in baseline_metrics:
|
||||
raise KeyError(
|
||||
f"Missing '{key}' in one of the metric dicts. "
|
||||
f"Available full={list(full_metrics)}, baseline={list(baseline_metrics)}"
|
||||
)
|
||||
"""R@1 gain from adding captions: full - baseline."""
|
||||
key = "r@1_query_to_gallery"
|
||||
return full_metrics[key] - baseline_metrics[key]
|
||||
|
||||
|
||||
@gin.configurable
|
||||
def run_evaluation_from_checkpoint(
|
||||
checkpoint_path: str,
|
||||
test_manifest: str,
|
||||
image_root: str,
|
||||
test_query_file: str,
|
||||
data_root: str,
|
||||
output_path: str = "eval_report.json",
|
||||
batch_size: int = 128,
|
||||
num_workers: int = 4,
|
||||
device: str = "cuda",
|
||||
) -> dict[str, float]:
|
||||
"""Standalone evaluation entry point (gin-configurable).
|
||||
|
||||
Args:
|
||||
checkpoint_path: Path to .pt checkpoint from training.
|
||||
test_manifest: Path to test manifest JSON.
|
||||
image_root: Directory prefix for images.
|
||||
output_path: Where to write the JSON report.
|
||||
batch_size: Batch size for encoding.
|
||||
num_workers: DataLoader workers.
|
||||
device: torch device.
|
||||
|
||||
Returns:
|
||||
Dict of retrieval metrics.
|
||||
"""
|
||||
"""Standalone evaluation from checkpoint."""
|
||||
from src.datasets.visloc_with_captions import (
|
||||
VisLocCaptionDataset,
|
||||
GeoLocCaptionDataset,
|
||||
collate_caption_batch,
|
||||
)
|
||||
|
||||
model = DualEncoderCaptionTest().to(device)
|
||||
ckpt = torch.load(checkpoint_path, map_location=device)
|
||||
ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
|
||||
model.load_state_dict(ckpt["model_state"])
|
||||
model.eval()
|
||||
|
||||
test_ds = VisLocCaptionDataset(
|
||||
manifest_path=test_manifest,
|
||||
image_root=image_root,
|
||||
test_ds = GeoLocCaptionDataset(
|
||||
query_file=test_query_file,
|
||||
data_root=data_root,
|
||||
image_transform=model.preprocess,
|
||||
)
|
||||
test_loader = DataLoader(
|
||||
@@ -222,15 +137,11 @@ def run_evaluation_from_checkpoint(
|
||||
pin_memory=True,
|
||||
)
|
||||
|
||||
metrics = evaluate_retrieval(
|
||||
model=model,
|
||||
loader=test_loader,
|
||||
device=device,
|
||||
)
|
||||
metrics = evaluate_retrieval(model=model, loader=test_loader, device=device)
|
||||
|
||||
report = {
|
||||
"checkpoint": checkpoint_path,
|
||||
"test_manifest": test_manifest,
|
||||
"test_query_file": test_query_file,
|
||||
"metrics": metrics,
|
||||
}
|
||||
out = Path(output_path)
|
||||
|
||||
Reference in New Issue
Block a user