A Transformer-Based Feature Segmentation and Region Alignment Method For UAV-View Geo-Localization

In plain words

A drone flies over a city and takes a photo of a building. To know where it is, the drone's photo can be matched against a large set of satellite images with known locations. The hard part is that the two photos look nothing alike: one is taken from above at low height and the other from far above, and the building may sit off-centre, be rotated or appear at a different scale.

The method, FSRA, first uses a transformer, a model that looks at all parts of an image at once. It cuts the image into small squares and asks which squares the model pays most attention to. It then groups the squares into a few bands, from the most active to the least active, and compares the matching bands from the drone photo and the satellite photo. Nobody labels these bands by hand; the grouping follows the model's own activity.

A second, simple trick is to show each satellite image to the model several times with small random changes. There is only one satellite image per building, so this gives the training more examples to learn from. It makes training slower but does not slow down the matching at use time. The authors report that this gives the best results at three repeats.

TL;DR

FSRA is a transformer-based method for matching drone photos to satellite images of the same building (University-1652 benchmark). It adds a patch-level heat segmentation (sort by mean activation, split into n regions) and a region-level alignment (average pooling per region, classification and triplet losses per region), on top of a ViT-S baseline with a global branch §3.1 §3.2. Satellite images are repeated k times during training §3.3. The reported numbers are from the authors and are above the CNN baseline LPN tab. 2.

Problem

Cross-view geo-localization must match an image from one viewpoint with the same place seen from another. Drone and satellite views differ in position, scale and rotation, and CNN features with pooling lose fine detail §1 fig. 1. The University-1652 benchmark has only one satellite image per building, which makes training unbalanced §3.3 §4.1 tab. 1.

Method

The method has three parts. The first is a ViT-S baseline that produces a global feature from the class token §3.1. The second is the heat segmentation (HSM) that splits patch tokens into n regions §3.2. The third is region alignment (HAB): every region is pooled and sent through its own classifier and triplet loss §3.2 fig. 5. A separate multiple-sampling strategy (k copies per satellite image) is applied during training §3.3.

1. Baseline. The input image is cut into N patches, projected linearly, and given a learnable position embedding and a learnable class token. A ViT encoder produces one token per patch and one for the class token. The class token is the global feature. The paper states the input sequence form and the learned position embedding §3.1 eq. 1.
2. Heat per patch. The heat of a patch is the mean of its feature values eq. 3.
3. Ordering and cutting. Patches are sorted from hottest to coolest. The sorted list is cut into n consecutive regions. Each region gets floor(N/n) patches, and the last region takes the remainder eq. 4 §3.2. The paper uses n = 3 §3.2 §4.4 fig. 8.
4. Region vectors. Each region is average-pooled into one vector eq. 5.
5. Heads and losses. Each region vector and the global feature go through a classifier head (linear layer, batch norm, dropout; the paper's figure 2 also lists a ReLU, see Paper vs code). They are trained with cross-entropy (ID loss, no label smoothing) and with the triplet loss, which is applied to every branch §3.1 §3.2 §3.4 fig. 2 fig. 3.
6. Triplet loss. Only pairs from different views are used. Margin 0.3 §3.4 eq. 10 eq. 11.
7. Multiple sampling. For each satellite image, k augmented copies are made, and k drone images of the same building are drawn at random §3.3. The paper reports k = 3 as best at batch size 8 §3.3 fig. 11 app. B.
8. Mutual learning (not in the final model). The paper defines a symmetric KL loss between the drone and satellite outputs eq. 6 eq. 7 eq. 8 eq. 9. The authors say it was not used in the final model because it did not add to the triplet loss §3.4 tab. 7.
9. Test. Distances are Euclidean §4.2 eq. 11. The paper does not say how the global and region vectors are combined for retrieval. See Paper vs code.

# Region features for one view  [§3.1] [§3.2]
# Input:  images (B, 3, 256, 256); n = number of regions (3 in the paper)
# Output: region_vectors: n tensors of shape (B, D); global_vector: (B, D)

def encode_view(images, n):
    # ViT tokens: the class token first, then N patch tokens  [§3.1] [eq. 1]
    # [code: models/FSRA/backbones/vit_pytorch.py#L352-L374]
    tokens = vit(images)                          # (B, 1 + N, D)
    global_vector = tokens[:, 0]                  # (B, D)
    patches = tokens[:, 1:]                       # (B, N, D)

    # Heat of a patch is the mean of its D feature values  [eq. 3]
    # [code: models/FSRA/make_model.py#L148]
    heat = mean(patches, over_channels)           # (B, N)

    # Order patches from hottest to coolest  [§3.2]
    # [code: models/FSRA/make_model.py#L150-L151]
    order = argsort_descending(heat)              # (B, N)
    ordered = gather(patches, order)              # (B, N, D)

    # Cut the ordered list into n regions  [eq. 4]
    # [code: models/FSRA/make_model.py#L154-L157]
    size = floor(N / n)
    sizes = [size] * (n - 1) + [N - (n - 1) * size]
    regions = split(ordered, sizes)

    # Average-pool each region into one vector  [eq. 5]
    # [code: models/FSRA/make_model.py#L159-L160]
    region_vectors = [mean(r, over_patches) for r in regions]
    return region_vectors, global_vector

# Loss for one batch of paired satellite and drone images  [§3.2] [§3.4]
# Input:  sat_logits, drone_logits: one (B, classes) tensor per branch
#         sat_feats, drone_feats:   one (B, 512) tensor per branch
#         labels: (B,) building ids, shared by both views
# Output: scalar loss

def batch_loss(sat_logits, drone_logits, sat_feats, drone_feats, labels):
    # ID loss: cross-entropy per branch, averaged over branches, then the
    # two views are summed
    # [code: losses/cal_loss.py#L5-L13] [code: train.py#L151]
    id_loss = mean_over_branches(cross_entropy(sat_logits, labels))
    id_loss = id_loss + mean_over_branches(cross_entropy(drone_logits, labels))

    # Triplet loss between views only, on the 512-d features of every branch
    # [eq. 10] [eq. 11]
    # [code: losses/cal_loss.py#L27-L39]
    # [code: losses/triplet_loss.py#L138-L186]
    triplet = cross_view_triplet(sat_feats, drone_feats, labels)
    triplet = mean_over_branches(triplet)

    # Unweighted sum; the KL term is off in the reported model  [§3.4]
    # [code: train.py#L169]
    return id_loss + triplet

Hyperparameters

Values in the paper are cited to their section. Items marked "code" come only from the repository at code_commit.

Results

All numbers below are author-reported. Setup: University-1652, drone and satellite views only, with a training set of 701 buildings and a test set of 951 buildings §4.1 tab. 1 §4.3. Drone→Satellite is the drone view target localization task; Satellite→Drone is drone navigation §4.1. Baseline: LPN with ResNet-50 tab. 2.

Dataset statistics tab. 1:

Comparison with other methods on University-1652 tab. 2:

Single-branch backbone comparison tab. 3. The loss setting for this table is not stated:

Padding robustness, AP (%) tab. 4. Black pad and flip pad shift the image edge by P pixels §4.4:

Input image size, R@1 and AP tab. 5. The second column group is labelled "Drone→Satellite" in the table. Its 256 px values equal the Satellite→Drone row of Table 2, so the second group is probably Satellite→Drone §4.4 tab. 2:

Distance between drone and building tab. 6. Drone→Satellite only:

Other tricks, AP (%) tab. 7:

Ground-view matching, appendix A tab. 8. The models are trained on ground, drone and satellite views and use only the classification loss app. A:

The paper says that ground-view R@1 "still remains in single digits" app. A.

Batch size 16 fig. 13: the paper reports that the best k is about 5 for batch size 16 and 3 for batch size 8. Values are only in the figure app. B.

Paper vs code

The official repository https://github.com/Dmmm1997/FSRA is linked in the paper's abstract, which satisfies the rule for an official repository. The owner name is not an author's name, so the link in the paper is the only evidence of ownership. Code citations are relative to that repository at code_commit. The repository has no LICENSE file, so no code is quoted verbatim. The first search step (links in the paper) found the repository; no third-party code was used.

1. Sampling. The paper samples k augmented satellite images per step §3.3. The code repeats each class index k times in a row, so a batch of 8 with k = 3 holds about three buildings, each two or three times Dataloader_University.py. The paper does not say how batches are formed.
2. Shipped run script. The script sets k = 1 and block = 3 train_test_local.sh. The train.py defaults are k = 1 and block = 1 train.py train.py. The k = 3 rows in Table 2 are therefore not what the script trains.
3. n = 1. The paper says the n = 1 setting adds global average pooling of the features to the global branch §4.4. The code returns only the global branch when block = 1 make_model.py.
4. n = 0. The paper says n = 0 uses only the global branch §4.4. The code has no such path: block = 0 divides by zero when it sizes the regions make_model.py (inferred from reading, not run).
5. KL loss. Eq. 9 sums both KL directions eq. 9. The code computes one direction per call, with log-softmax of one view and softmax of the other cal_loss.py train.py. It is off by default train.py. Eq. 7 defines p as a log-softmax, which conflicts with the form of eq. 6; the code follows the usual convention.
6. Triplet features. Eq. 10 uses the raw distance between features eq. 10 eq. 11. The code L2-normalises the features before measuring distances triplet_loss.py triplet_loss.py. The paper does not say this.
7. Hard mining. The paper says nothing about how positives and negatives are chosen. The code uses the hardest cross-view positive and the hardest cross-view negative for each anchor triplet_loss.py.
8. Which branches get the triplet loss. The paper applies the triplet loss to all regions §3.2. The code applies it to every branch, including the global one, and averages over branches train.py cal_loss.py. The triplet is computed on the 512-d bottleneck features, not the 768-d ViT output make_model.py.
9. ReLU in the head. Figure 2 lists a ReLU in the classifier fig. 2. The code uses no ReLU by default, because the head is built without the relu option make_model.py make_model.py make_model.py. LeakyReLU(0.1) is added only when the option is set make_model.py.
10. Classifier init. The paper says Kaiming init for the classifier §4.2. The code applies Kaiming init to the 768→512 linear layer and N(0, 0.001) to the final layer make_model.py make_model.py.
11. Each region has its own head. The code builds one head per region make_model.py. The paper does not say whether heads are shared.
12. Final norm. The code applies the ViT's final LayerNorm before computing heat and pooling vit_pytorch.py. The paper says only "the output of the transformer" §3.2.
13. Loss aggregation. Cross-entropy is averaged over branches, then summed over views cal_loss.py train.py. The total is an unweighted sum train.py. The paper states neither.
14. Satellite augmentation. The code rotates satellite images by up to 90 degrees make_dataloader.py. The paper does not list rotation §3.3. Colour jitter is applied only with a flag that is off by default train.py make_dataloader.py.
15. Backbone depth. The paper gives no depth. The code builds 8 transformer blocks with 8 heads and an MLP ratio of 3 vit_pytorch.py. The script loads a file named vit_small_p16_224-15ec54c9.pth train_test_local.sh. The repository ships no weights (pretrain_model/ is empty). Whether that file fits 8 blocks was not verified.
16. Position embedding. The paper says the position embedding cannot be loaded from ImageNet and is learned §3.1. The code's loader resizes the pretrained position embedding to the new grid and copies it vit_pytorch.py.
17. Test-time features. The code concatenates the global vector and all region vectors, L2-normalises each part, and flattens them make_model.py test_server.py. Each part then has the same norm, so ranking by dot product is the same as ranking by Euclidean distance (derived, not run). The paper does not state how region vectors are combined.
18. Test-time flip. The code sums features of the image and its horizontal flip before normalising test_server.py. The paper does not mention it.
19. Ranking order. The code sorts scores and reverses the order so the highest similarity comes first evaluate_gpu.py.
20. Not in the paper. Nesterov momentum make_optimizer.py and mixed precision by default train.py train.py.
21. Partial batch. The code skips a last batch smaller than the batch size train.py. The paper does not say.
22. Table 2 and Table 7 disagree. The k = 1 row of Table 2 (84.82 / 81.53) matches the KLLoss-plus-triplet row of Table 7, but §3.4 says KLLoss was not used. The triplet-only row is 84.85 / 81.52 tab. 2 tab. 7 §3.4. The paper does not say which one is reported.
23. Evaluation script. load_network calls the model constructors without the opt argument utils_server.py model.py. As committed, test_server.py should fail on loading a checkpoint. This was inferred from reading the code and not run, because no data or GPU was available.

Pitfalls

• Heat is a mean over channels, not a learned score. The region boundaries come from the sorted order of mean activations eq. 3 make_model.py. Ties and the choice of the final norm (point 12) change which patches land in which region. The paper says the regions correspond to buildings, roads and trees §4.4 fig. 4. No segmentation metric is reported, so this reading is an interpretation of heat bands.
• Region sizes for N = 256. With a 256 px input and patch size 16, N is 256 tab. 3 §4.2 (derived). Eq. 4 then gives 85, 85 and 86 patches. The last region is the largest. At 512 px, N = 1024 (derived).
• Batch composition with k > 1. Repeats of one building sit next to each other in the sampler, so a batch of 8 at k = 3 has few negatives for the triplet loss Dataloader_University.py. The paper attributes the k = 3 gain to the balance of positives and negatives §4.4. The effect of the adjacent repeats was not measured.
• Cost. k changes training time only, not inference §3.3 §4.4. The paper says training time grows with k app. B. It gives no training-time numbers. ViT-S inference is 1.21× ResNet-50 tab. 3.
• Triplet margin is on normalised features. With normalisation (point 6), distances lie on the unit sphere, so margin 0.3 means something different from the same margin on raw features.
• Position embedding at other sizes. The paper says the position embedding is learned and not loaded from ImageNet §3.1. The code resizes it when the grid changes (point 16). Which path the 512 px run took is not stated.
• Black pad at test time. The code's query transform pads the left edge with zeros and crops the right edge queryDataset.py. It is applied to the query images only test_server.py. The padding ablation tab. 4 is not given as a script.
• Batch-size dependence of k. The paper reports k = 3 at batch 8 and k = 5 at batch 16 app. B fig. 13. Choose k with the batch size.

Lineage

• Builds on: ViT for the backbone §3.1; the part-based idea from LPN and related part methods §2.1; the equal-multiplicity sampling idea from LCM §1 §3.3.
• Compared against: LPN, LCM, instance loss with GeM pooling and triplet and contrastive losses on VGG16 and ResNet-50 tab. 2.
• Related transformer work named by the authors: TransReID for retrieval, and DETR and SETR for other vision tasks §2.2.