Remove projections (1024 native), add satellite text, dual GatedFusion
Architecture changes: - Removed proj_drone/proj_sat (1024→512): retrieval space is now DINOv3 native 1024-dim, no information loss from projection - TextFusionMLP: 2304→1024→1024 (was 2304→768→512), shared between branches - Gallery branch now uses satellite captions (L1/L2/L3) via shared TextFusionMLP - Two separate GatedFusion gates: α_q (query) and α_g (gallery) - For sat images without captions (~57%): gate passes image features through Dataset changes: - GTAUAVDataset now loads satellite captions from caption index - collate_gtauav_batch includes sat_caption_l1/l2/l3 Training loop: - Passes satellite captions to model forward - Logs both gate_q and gate_g values 11.1M trainable / 734M total (1.51%) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -176,13 +176,20 @@ class GTAUAVDataset(Dataset):
|
||||
else:
|
||||
continue # No match, skip.
|
||||
|
||||
# Get captions.
|
||||
# Get drone captions.
|
||||
cap_data = self.caption_index.get(drone_name)
|
||||
if cap_data is not None:
|
||||
l1, l2, l3 = _parse_caption_levels(cap_data["output"])
|
||||
else:
|
||||
l1 = l2 = l3 = _EMPTY_CAPTION
|
||||
|
||||
# Pre-parse satellite captions for all candidates.
|
||||
sat_captions: dict[str, tuple[str, str, str]] = {}
|
||||
for sat_name in sat_candidates:
|
||||
sat_cap = self.caption_index.get(sat_name)
|
||||
if sat_cap is not None:
|
||||
sat_captions[sat_name] = _parse_caption_levels(sat_cap["output"])
|
||||
|
||||
self.entries.append({
|
||||
"drone_name": drone_name,
|
||||
"drone_dir": pair["drone_img_dir"],
|
||||
@@ -192,6 +199,7 @@ class GTAUAVDataset(Dataset):
|
||||
"caption_l1": l1,
|
||||
"caption_l2": l2,
|
||||
"caption_l3": l3,
|
||||
"sat_captions": sat_captions,
|
||||
})
|
||||
|
||||
def _load_image(self, directory: str, filename: str, transform: Callable | None = None) -> torch.Tensor:
|
||||
@@ -222,18 +230,28 @@ class GTAUAVDataset(Dataset):
|
||||
|
||||
sat_img = self._load_image(entry["sat_dir"], sat_name, self.sat_transform)
|
||||
|
||||
# Captions with optional dropout.
|
||||
# Drone captions with optional dropout.
|
||||
if self.drop_caption_prob > 0 and self._rng.random() < self.drop_caption_prob:
|
||||
l1 = l2 = l3 = _EMPTY_CAPTION
|
||||
else:
|
||||
l1, l2, l3 = entry["caption_l1"], entry["caption_l2"], entry["caption_l3"]
|
||||
|
||||
# Satellite captions (empty string if not available).
|
||||
sat_caps = entry["sat_captions"].get(sat_name)
|
||||
if sat_caps is not None:
|
||||
sat_l1, sat_l2, sat_l3 = sat_caps
|
||||
else:
|
||||
sat_l1 = sat_l2 = sat_l3 = _EMPTY_CAPTION
|
||||
|
||||
return {
|
||||
"drone_img": drone_img,
|
||||
"sat_img": sat_img,
|
||||
"caption_l1": l1,
|
||||
"caption_l2": l2,
|
||||
"caption_l3": l3,
|
||||
"sat_caption_l1": sat_l1,
|
||||
"sat_caption_l2": sat_l2,
|
||||
"sat_caption_l3": sat_l3,
|
||||
"pair_id": entry["drone_name"],
|
||||
}
|
||||
|
||||
@@ -248,5 +266,8 @@ def collate_gtauav_batch(
|
||||
"caption_l1": [b["caption_l1"] for b in batch],
|
||||
"caption_l2": [b["caption_l2"] for b in batch],
|
||||
"caption_l3": [b["caption_l3"] for b in batch],
|
||||
"sat_caption_l1": [b["sat_caption_l1"] for b in batch],
|
||||
"sat_caption_l2": [b["sat_caption_l2"] for b in batch],
|
||||
"sat_caption_l3": [b["sat_caption_l3"] for b in batch],
|
||||
"pair_ids": [b["pair_id"] for b in batch],
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user