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>
This commit is contained in:
2
GeoLoc-UAV-main/models/backbone/__init__.py
Normal file
2
GeoLoc-UAV-main/models/backbone/__init__.py
Normal file
@@ -0,0 +1,2 @@
|
||||
from .resnet import ResNet
|
||||
from .dinov2 import DINOv2
|
||||
94
GeoLoc-UAV-main/models/backbone/dinov2.py
Normal file
94
GeoLoc-UAV-main/models/backbone/dinov2.py
Normal file
@@ -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
|
||||
107
GeoLoc-UAV-main/models/backbone/resnet.py
Normal file
107
GeoLoc-UAV-main/models/backbone/resnet.py
Normal file
@@ -0,0 +1,107 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision
|
||||
import numpy as np
|
||||
|
||||
class ResNet(nn.Module):
|
||||
def __init__(self,
|
||||
model_name='resnet50',
|
||||
pretrained=True,
|
||||
layers_to_freeze=2,
|
||||
layers_to_crop=[],
|
||||
pretrain_flag = False
|
||||
):
|
||||
"""Class representing the resnet backbone used in the pipeline
|
||||
we consider resnet network as a list of 5 blocks (from 0 to 4),
|
||||
layer 0 is the first conv+bn and the other layers (1 to 4) are the rest of the residual blocks
|
||||
we don't take into account the global pooling and the last fc
|
||||
|
||||
Args:
|
||||
model_name (str, optional): The architecture of the resnet backbone to instanciate. Defaults to 'resnet50'.
|
||||
pretrained (bool, optional): Whether pretrained or not. Defaults to True.
|
||||
layers_to_freeze (int, optional): The number of residual blocks to freeze (starting from 0) . Defaults to 2.
|
||||
layers_to_crop (list, optional): Which residual layers to crop, for example [3,4] will crop the third and fourth res blocks. Defaults to [].
|
||||
|
||||
Raises:
|
||||
NotImplementedError: if the model_name corresponds to an unknown architecture.
|
||||
"""
|
||||
super().__init__()
|
||||
self.model_name = model_name.lower()
|
||||
self.layers_to_freeze = layers_to_freeze
|
||||
self.flag = pretrain_flag
|
||||
|
||||
if pretrained:
|
||||
# the new naming of pretrained weights, you can change to V2 if desired.
|
||||
weights = 'IMAGENET1K_V1'
|
||||
else:
|
||||
weights = None
|
||||
|
||||
if 'swsl' in model_name or 'ssl' in model_name:
|
||||
# These are the semi supervised and weakly semi supervised weights from Facebook
|
||||
self.model = torch.hub.load(
|
||||
'facebookresearch/semi-supervised-ImageNet1K-models', model_name)
|
||||
else:
|
||||
if 'resnext50' in model_name:
|
||||
self.model = torchvision.models.resnext50_32x4d(
|
||||
weights=weights)
|
||||
elif 'resnet50' in model_name:
|
||||
self.model = torchvision.models.resnet50(weights=weights)
|
||||
elif '101' in model_name:
|
||||
self.model = torchvision.models.resnet101(weights=weights)
|
||||
elif '152' in model_name:
|
||||
self.model = torchvision.models.resnet152(weights=weights)
|
||||
elif '34' in model_name:
|
||||
self.model = torchvision.models.resnet34(weights=weights)
|
||||
elif '18' in model_name:
|
||||
# self.model = torchvision.models.resnet18(pretrained=False)
|
||||
self.model = torchvision.models.resnet18(weights=weights)
|
||||
elif 'wide_resnet50_2' in model_name:
|
||||
self.model = torchvision.models.wide_resnet50_2(
|
||||
weights=weights)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
'Backbone architecture not recognized!')
|
||||
|
||||
# freeze only if the model is pretrained
|
||||
if pretrained and self.flag:
|
||||
if layers_to_freeze >= 0:
|
||||
self.model.conv1.requires_grad_(False)
|
||||
self.model.bn1.requires_grad_(False)
|
||||
if layers_to_freeze >= 1:
|
||||
self.model.layer1.requires_grad_(False)
|
||||
if layers_to_freeze >= 2:
|
||||
self.model.layer2.requires_grad_(False)
|
||||
if layers_to_freeze >= 3:
|
||||
self.model.layer3.requires_grad_(False)
|
||||
|
||||
# remove the avgpool and most importantly the fc layer
|
||||
self.model.avgpool = None
|
||||
self.model.fc = None
|
||||
|
||||
if 4 in layers_to_crop:
|
||||
self.model.layer4 = None
|
||||
if 3 in layers_to_crop:
|
||||
self.model.layer3 = None
|
||||
|
||||
out_channels = 2048
|
||||
if '34' in model_name or '18' in model_name:
|
||||
out_channels = 256
|
||||
|
||||
self.out_channels = out_channels // 2 if self.model.layer4 is None else out_channels
|
||||
self.out_channels = self.out_channels // 2 if self.model.layer3 is None else self.out_channels
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
x = self.model.conv1(x)
|
||||
x = self.model.bn1(x)
|
||||
x = self.model.relu(x)
|
||||
x = self.model.maxpool(x)
|
||||
x = self.model.layer1(x)
|
||||
x = self.model.layer2(x)
|
||||
if self.model.layer3 is not None:
|
||||
x = self.model.layer3(x)
|
||||
if self.model.layer4 is not None:
|
||||
x = self.model.layer4(x)
|
||||
|
||||
return x
|
||||
|
||||
Reference in New Issue
Block a user