Initial import: World-UAV prepro

Add dataloaders (v1/v2), analysis scripts, and documentation for working with UAV-GeoLoc (World-UAV).

Co-authored-by: Cursor <cursoragent@cursor.com>

GeoLoc-UAV-main/models/aggregators/LPN.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn as nn
+import math
+from torch.nn import functional as F
+
+def get_part_pool(x, block=4, no_overlap=True):
+        result = []
+        H, W = x.size(2), x.size(3)
+        c_h, c_w = int(H/2), int(W/2)
+        per_h, per_w = H/(2*block),W/(2*block)
+        if per_h < 1 and per_w < 1:
+            new_H, new_W = H+(block-c_h)*2, W+(block-c_w)*2
+            x = nn.functional.interpolate(x, size=[new_H,new_W], mode='bilinear', align_corners=True)
+            H, W = x.size(2), x.size(3)
+            c_h, c_w = int(H/2), int(W/2)
+            per_h, per_w = H/(2*block),W/(2*block)
+        per_h, per_w = math.floor(per_h), math.floor(per_w)
+        for i in range(block):
+            i = i + 1
+            if i < block:
+                x_curr = x[:,:,(c_h-i*per_h):(c_h+i*per_h),(c_w-i*per_w):(c_w+i*per_w)]
+                if no_overlap and i > 1:
+                    x_pre = x[:,:,(c_h-(i-1)*per_h):(c_h+(i-1)*per_h),(c_w-(i-1)*per_w):(c_w+(i-1)*per_w)] 
+                    x_pad = F.pad(x_pre,(per_h,per_h,per_w,per_w),"constant",0)
+                    x_curr = x_curr - x_pad
+                result.append(x_curr)
+            else:
+                if no_overlap and i > 1:
+                    x_pre = x[:,:,(c_h-(i-1)*per_h):(c_h+(i-1)*per_h),(c_w-(i-1)*per_w):(c_w+(i-1)*per_w)]
+                    pad_h = c_h-(i-1)*per_h
+                    pad_w = c_w-(i-1)*per_w
+                    # x_pad = F.pad(x_pre,(pad_h,pad_h,pad_w,pad_w),"constant",0)
+                    if x_pre.size(2)+2*pad_h == H:
+                        x_pad = F.pad(x_pre,(pad_h,pad_h,pad_w,pad_w),"constant",0)
+                    else:
+                        ep = H - (x_pre.size(2)+2*pad_h)
+                        x_pad = F.pad(x_pre,(pad_h+ep,pad_h,pad_w+ep,pad_w),"constant",0)
+                    x = x - x_pad
+                result.append(x_curr)
+        return result

GeoLoc-UAV-main/models/aggregators/__init__.py

@@ -0,0 +1,3 @@
+from .gem import GeMPool
+from .convap import ConvAP
+from .multiconvap import MulConvAP

GeoLoc-UAV-main/models/aggregators/convap.py

@@ -0,0 +1,33 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+
+
+class ConvAP(nn.Module):
+    """Implementation of ConvAP as of https://arxiv.org/pdf/2210.10239.pdf
+
+    Args:
+        in_channels (int): number of channels in the input of ConvAP
+        out_channels (int, optional): number of channels that ConvAP outputs. Defaults to 512.
+        s1 (int, optional): spatial height of the adaptive average pooling. Defaults to 2.
+        s2 (int, optional): spatial width of the adaptive average pooling. Defaults to 2.
+    """
+    def __init__(self, in_channels, out_channels=512, s1=2, s2=2):
+        super(ConvAP, self).__init__()
+        self.channel_pool = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=True)
+        self.AAP = nn.AdaptiveAvgPool2d((s1, s2))
+
+    def forward(self, x):
+        #
+        x, t = x #dinov2专属
+        # x = self.channel_pool(x)
+        x = self.AAP(x)
+        x = F.normalize(x.flatten(1), p=2, dim=1)
+        return x
+    
+
+if __name__ == '__main__':
+    x = torch.randn(4, 2048, 10, 10)
+    m = ConvAP(2048, 512)
+    r = m(x)
+    print(r.shape)

GeoLoc-UAV-main/models/aggregators/gem.py

@@ -0,0 +1,17 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+
+class GeMPool(nn.Module):
+    """Implementation of GeM as in https://github.com/filipradenovic/cnnimageretrieval-pytorch
+    we add flatten and norm so that we can use it as one aggregation layer.
+    """
+    def __init__(self, p=3, eps=1e-6):
+        super().__init__()
+        self.p = nn.Parameter(torch.ones(1)*p)
+        self.eps = eps
+
+    def forward(self, x):
+        x = F.avg_pool2d(x.clamp(min=self.eps).pow(self.p), (x.size(-2), x.size(-1))).pow(1./self.p)
+        x = x.flatten(1)
+        return F.normalize(x, p=2, dim=1)

GeoLoc-UAV-main/models/aggregators/multiconvap.py

@@ -0,0 +1,96 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+
+from models.aggregators.LPN import get_part_pool
+
+class L2Norm(nn.Module):
+    def __init__(self, dim=1):
+        super().__init__()
+        self.dim = dim
+    def forward(self, x):
+        return F.normalize(x, p=2, dim=self.dim)
+
+class GeMPool(nn.Module):
+    """Implementation of GeM as in https://github.com/filipradenovic/cnnimageretrieval-pytorch
+    we add flatten and norm so that we can use it as one aggregation layer.
+    """
+    def __init__(self, p=3, eps=1e-6):
+        super().__init__()
+        self.p = nn.Parameter(torch.ones(1)*p)
+        self.eps = eps
+
+    def forward(self, x):
+        x = F.avg_pool2d(x.clamp(min=self.eps).pow(self.p), (x.size(-2), x.size(-1))).pow(1./self.p)
+        x = x.flatten(1)
+        return F.normalize(x, p=2, dim=1)
+
+class MulConvAP(nn.Module):
+    """Implementation of ConvAP as of https://arxiv.org/pdf/2210.10239.pdf
+
+    Args:
+        in_channels (int): number of channels in the input of ConvAP
+        out_channels (int, optional): number of channels that ConvAP outputs. Defaults to 512.
+        s1 (int, optional): spatial height of the adaptive average pooling. Defaults to 2.
+        s2 (int, optional): spatial width of the adaptive average pooling. Defaults to 2.
+    """
+    def __init__(self, in_channels, out_channels=512, s1=2, s2=2, LPN=False):
+        super(MulConvAP, self).__init__()
+        self.out_channels = out_channels
+        self.channel_pool_1 = nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=1, bias=True)
+        self.channel_pool_3 = nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=3, padding=1,bias=True)
+        self.channel_pool_5 = nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=5, padding=2,bias=True)
+      
+        # self.AAP = nn.AdaptiveAvgPool2d((s1, s2))
+        self.AAP = nn.Sequential(L2Norm(), GeMPool())
+
+        # using LPN
+        if LPN == True:
+            self.LPN = True
+        else:
+            self.LPN = False
+    def forward(self, x):
+
+        if self.LPN == False:
+            # x, t = x #dinov2专属
+            x1 = self.channel_pool_1(x)
+            x3 = self.channel_pool_3(x)
+            x5 = self.channel_pool_5(x)
+
+            x1 = self.AAP(x1)
+            x3 = self.AAP(x3)
+            x5 = self.AAP(x5)
+
+            x = [i for i in [x1, x3, x5]]
+            x = torch.cat(x,dim=1)
+
+            # x = self.AAP(x)
+            x = F.normalize(x.flatten(1), p=2, dim=1)
+            return x
+        else:
+            partition_feature = get_part_pool(x)
+            partition_feature_list = []
+            for one_feature in partition_feature:
+                x1 = self.channel_pool_1(one_feature)
+                x3 = self.channel_pool_3(one_feature)
+                x5 = self.channel_pool_5(one_feature)
+
+                x1 = self.AAP(x1)
+                x3 = self.AAP(x3)
+                x5 = self.AAP(x5)
+
+                x = [i for i in [x1, x3, x5]]
+                x = torch.cat(x,dim=1)
+
+                x = F.normalize(x.flatten(1), p=2, dim=1)
+                partition_feature_list.append(x)
+            # partition_feature_tensor = torch.stack(partition_feature_list, dim=2).reshape(x.shape[0], -1)
+            
+            return partition_feature_list
+
+
+if __name__ == '__main__':
+    x = torch.randn(4, 2048, 10, 10)
+    # m = ConvAP(2048, 512)
+    # r = m(x)
+    # print(r.shape)