Pikaliov
4ff36ce188
Add dataloaders (v1/v2), analysis scripts, and documentation for working with UAV-GeoLoc (World-UAV). Co-authored-by: Cursor <cursoragent@cursor.com>
41 lines
1.8 KiB
Python
41 lines
1.8 KiB
Python
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
|