Pikaliov/World-UAV-ds
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial import: World-UAV prepro

Browse Source 
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:
Pikaliov
2026-05-09 12:44:49 +03:00
commit 4ff36ce188
72 changed files with 13594 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

167
GeoLoc-UAV-main/models/game4loc.py Normal file
View File

@@ -0,0 +1,167 @@
import torch
import timm
import numpy as np
import torch.nn as nn
from PIL import Image
from urllib.request import urlopen
from thop import profile
class MLP(nn.Module):
def __init__(self, input_size=2048, hidden_size=512, output_size=2):
super(MLP, self).__init__()
self.fc1 = nn.Linear(input_size, hidden_size)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_size, hidden_size // 2)
self.fc3 = nn.Linear(hidden_size // 2, output_size)
def forward(self, x):
x = self.fc1(x)
x = self.relu(x)
x = self.fc2(x)
x = self.relu(x)
x = self.fc3(x)
return x
class DesModel(nn.Module):
def __init__(self,
model_name='vit',
pretrained=True,
img_size=384,
share_weights=True,
train_with_recon=False,
train_with_offset=False,
model_hub='timm'):
super(DesModel, self).__init__()
self.share_weights = share_weights
self.model_name = model_name
self.img_size = img_size
if share_weights:
if "vit" in model_name or "swin" in model_name:
# automatically change interpolate pos-encoding to img_size
self.model = timm.create_model(model_name, pretrained=pretrained, num_classes=0, img_size=img_size)
else:
self.model = timm.create_model(model_name, pretrained=pretrained, num_classes=0)
else:
if "vit" in model_name or "swin" in model_name:
self.model1 = timm.create_model(model_name, pretrained=pretrained, num_classes=0, img_size=img_size)
self.model2 = timm.create_model(model_name, pretrained=pretrained, num_classes=0, img_size=img_size)
else:
self.model1 = timm.create_model(model_name, pretrained=pretrained, num_classes=0)
self.model2 = timm.create_model(model_name, pretrained=pretrained, num_classes=0)
if train_with_offset:
self.MLP = MLP()
self.logit_scale = torch.nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
def get_config(self,):
if self.share_weights:
data_config = timm.data.resolve_model_data_config(self.model)
else:
data_config = timm.data.resolve_model_data_config(self.model1)
return data_config
def set_grad_checkpointing(self, enable=True):
if self.share_weights:
self.model.set_grad_checkpointing(enable)
else:
self.model1.set_grad_checkpointing(enable)
self.model2.set_grad_checkpointing(enable)
def freeze_layers(self, frozen_blocks=10, frozen_stages=[0,0,0,0]):
pass
def forward(self, img1=None, img2=None):
if self.share_weights:
if img1 is not None and img2 is not None:
image_features1 = self.model(img1)
image_features2 = self.model(img2)
return image_features1, image_features2
elif img1 is not None:
image_features = self.model(img1)
return image_features
else:
image_features = self.model(img2)
return image_features
else:
if img1 is not None and img2 is not None:
image_features1 = self.model1(img1)
image_features2 = self.model2(img2)
return image_features1, image_features2
elif img1 is not None:
image_features = self.model1(img1)
return image_features
else:
image_features = self.model2(img2)
return image_features
def offset_pred(self, img_feature1, img_feature2):
offset = self.MLP(torch.cat((img_feature1, img_feature2), dim=1))
return offset
if __name__ == '__main__':
# model = TimmModel(model_name='timm/vit_large_patch16_384.augreg_in21k_ft_in1k')
# # model = TimmModel(model_name='timm/vit_base_patch16_224.augreg_in1k')
# # from timm.models.vision_transformer import vit_base_patch16_224
# # model = vit_base_patch16_224(img_size=384, patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, num_classes=0)
# model = DesModel(model_name='timm/resnet101.tv_in1k', img_size=384)
# model = DesModel(model_name='convnext_base.fb_in22k_ft_in1k_384', img_size=384)
model = DesModel(model_name='timm/swin_base_patch4_window7_224.ms_in22k_ft_in1k', img_size=384)
# # model = TimmModel(model_name='vit_base_patch16_rope_reg1_gap_256.sbb_in1k')
# # model = TimmModel(model_name='timm/vit_medium_patch16_rope_reg1_gap_256.sbb_in1k')
# # model = TimmModel(model_name='timm/vit_medium_patch16_gap_256.sw_in12k_ft_in1k')
# # model = TimmModel(model_name='timm/resnet101.tv_in1k')
# # img = Image.open(urlopen(
# # 'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png'
# # ))
x = torch.rand((1, 3, 384, 384))
x = x.cuda()
model.cuda()
x = model(x)
print(x.shape)
# flops, params = profile(model, inputs=(x,))
# # print(img.size)
# # img = transform(img)
# # print(img.size)
# # print(model1)
# print('flops(G)', flops/1e9, 'params(M)', params/1e6)
# from transformers import CLIPProcessor, CLIPModel
# model = CLIPModel.from_pretrained("/home/xmuairmud/jyx/clip-vit-base-patch16")
# vision_model = model.vision_model
# print(vision_model)
# dinov2_vitb14_reg = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14_reg')
# print(dinov2_vitb14_reg.set_grad_checkpointing(True))
# from transformers import ViTModel, ViTImageProcessor, AutoModelForImageClassification, AutoConfig
# config = AutoConfig.from_pretrained('facebook/dino-vitb16')
# config.image_size = 384
# model = ViTModel.from_pretrained('facebook/dino-vitb16', config=config, ignore_mismatched_sizes=True)
# model = timm.create_model('vit_base_patch14_reg4_dinov2.lvd142m', pretrained=True, img_size=(384, 384))
# data_config = timm.data.resolve_model_data_config(model)
# print(data_config)
# processor = ViTImageProcessor.from_pretrained('facebook/dino-vitb16')
# x = torch.rand((1, 3, 384, 384))
# inputs = processor(images=x, return_tensors="pt")
# print(inputs['pixel_values'].shape)
# outputs = model(**inputs)
# print(outputs.pooler_output.shape)
# print(model(x).shape)
# flops, params = profile(dinov2_vitb14_reg, inputs=(x,))
# print('flops(G)', flops/1e9, 'params(M)', params/1e6)