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/backbone/dinov2.py

@@ -0,0 +1,94 @@
+import torch
+import torch.nn as nn
+
+DINOV2_ARCHS = {
+    'dinov2_vits14': 384,
+    'dinov2_vitb14': 768,
+    'dinov2_vitl14': 1024,
+    'dinov2_vitg14': 1536,
+}
+
+class DINOv2(nn.Module):
+    """
+    DINOv2 model
+
+    Args:
+        model_name (str): The name of the model architecture 
+            should be one of ('dinov2_vits14', 'dinov2_vitb14', 'dinov2_vitl14', 'dinov2_vitg14')
+        num_trainable_blocks (int): The number of last blocks in the model that are trainable.
+        norm_layer (bool): If True, a normalization layer is applied in the forward pass.
+        return_token (bool): If True, the forward pass returns both the feature map and the token.
+    """
+    def __init__(
+            self,
+            model_name='dinov2_vitb14',
+            num_trainable_blocks=2,
+            norm_layer=False,
+            return_token=False,
+            pretrain_flag=False
+        ):
+        super().__init__()
+
+        assert model_name in DINOV2_ARCHS.keys(), f'Unknown model name {model_name}'
+        
+        self.model = torch.hub.load('facebookresearch/dinov2', "dinov2_vits14")
+        # torch.hub.load('/home/Shen/.cache/torch/hub/facebookresearch_dinov2_main/',
+                                    # model_name,
+            # source='local')
+
+        # self.model = torch.hub.load('facebookresearch/dinov2', model_name)
+        self.num_channels = DINOV2_ARCHS[model_name]
+        self.num_trainable_blocks = num_trainable_blocks
+        self.norm_layer = norm_layer
+        self.return_token = return_token
+        self.flag = pretrain_flag
+
+
+    def forward(self, x):
+        """
+        The forward method for the DINOv2 class
+
+        Parameters:
+            x (torch.Tensor): The input tensor [B, 3, H, W]. H and W should be divisible by 14.
+
+        Returns:
+            f (torch.Tensor): The feature map [B, C, H // 14, W // 14].
+            t (torch.Tensor): The token [B, C]. This is only returned if return_token is True.
+        """
+
+        B, C, H, W = x.shape
+
+        x = self.model.prepare_tokens_with_masks(x)
+        if self.flag:
+        # When flag is True, freeze all parameters
+            for param in self.model.parameters():
+                param.requires_grad = False
+            with torch.no_grad():
+                for blk in self.model.blocks:
+                    x = blk(x)
+        else:
+            # When flag is False, freeze part of the parameters (e.g., first blocks)
+            for param in self.model.parameters():
+                param.requires_grad = False  # Freeze all layers initially
+            # Unfreeze the last few blocks (trainable)
+            for param in self.model.blocks[-self.num_trainable_blocks:].parameters():
+                param.requires_grad = True
+
+            with torch.no_grad():
+                for blk in self.model.blocks[:-self.num_trainable_blocks]:  # Freeze these blocks
+                    x = blk(x)
+            # Last blocks are trained
+            for blk in self.model.blocks[-self.num_trainable_blocks:]:  # Train these blocks
+                x = blk(x)
+            if self.norm_layer:
+                x = self.model.norm(x)
+        
+        t = x[:, 0]
+        f = x[:, 1:]
+
+        # Reshape to (B, C, H, W)
+        f = f.reshape((B, H // 14, W // 14, self.num_channels)).permute(0, 3, 1, 2)
+
+        if self.return_token:
+            return f, t
+        return f