Pikaliov
4ff36ce188
Add dataloaders (v1/v2), analysis scripts, and documentation for working with UAV-GeoLoc (World-UAV). Co-authored-by: Cursor <cursoragent@cursor.com>
957 lines
27 KiB
Python
957 lines
27 KiB
Python
import os
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image, UnidentifiedImageError
|
|
from torch.utils.data import Dataset
|
|
import copy
|
|
from tqdm import tqdm
|
|
import time
|
|
import random
|
|
import glob
|
|
import json
|
|
import pandas as pd
|
|
|
|
from torch.utils.data import DataLoader
|
|
import torchvision.transforms as T
|
|
import json
|
|
|
|
|
|
from utils.utils import TransformerCV
|
|
|
|
# transform_config = {
|
|
# "sample_scale_begin": 0,
|
|
# "sample_scale_inter": 0.5,
|
|
# "sample_scale_num": 3,
|
|
# "sample_rotate_begin": 0,
|
|
# "sample_rotate_inter": 45,
|
|
# "sample_rotate_num": 8,
|
|
# }
|
|
json_path = "/media/guan/新加卷/Code/Code/configs/transform_config.json"
|
|
with open(json_path, 'r', encoding='utf-8') as file:
|
|
data = json.load(file)
|
|
transform_config = data["transform_config"]
|
|
|
|
|
|
# transform_config = {
|
|
# "sample_scale_begin": 0,
|
|
# "sample_scale_inter": 0.5,
|
|
# "sample_scale_num": 1,
|
|
# "sample_rotate_begin": 0,
|
|
# "sample_rotate_inter": 0,
|
|
# "sample_rotate_num": 1,
|
|
# }
|
|
|
|
default_transform = T.Compose([
|
|
T.ToTensor(),
|
|
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
])
|
|
|
|
def get_data(txt):
|
|
|
|
data = {}
|
|
idx = 0
|
|
with open(txt, 'r') as f:
|
|
for line in f:
|
|
line_list = line.split(' ')[:-1]
|
|
data[idx] = line_list
|
|
idx += 1
|
|
|
|
return data
|
|
|
|
class WorldDatasetTrainGroup(Dataset):
|
|
def __init__(self,
|
|
data_dir,
|
|
query_txt,
|
|
transforms_query=default_transform,
|
|
transforms_db=default_transform,
|
|
shuffle_batch_size=64):
|
|
super().__init__()
|
|
|
|
self.pairs = []
|
|
self.data = get_data(query_txt)
|
|
|
|
for idx in self.data.items():
|
|
query_img_path = os.path.join(data_dir, idx[1][0])
|
|
label = eval(idx[1][1])
|
|
db_image_path = os.path.join(data_dir, idx[1][2])
|
|
self.pairs.append((label, query_img_path, db_image_path))
|
|
|
|
self.transforms_query = transforms_query
|
|
self.transforms_db = transforms_db
|
|
self.shuffle_batch_size = shuffle_batch_size
|
|
|
|
self.samples = copy.deepcopy(self.pairs)
|
|
|
|
self.group_transformer = TransformerCV(transform_config)
|
|
self.pts_step = 5
|
|
|
|
def __getitem__(self, index):
|
|
|
|
idx, query_img_path, db_img_path = self.samples[index]
|
|
# query
|
|
query_img = self.image_loader(query_img_path)
|
|
# db
|
|
db_img = self.image_loader(db_img_path)
|
|
# image transforms
|
|
if self.transforms_query is not None:
|
|
query_img = self.transforms_query(query_img)
|
|
|
|
if self.transforms_db is not None:
|
|
db_img = self.transforms_db(db_img)
|
|
|
|
# return query_img, db_img, idx
|
|
# group
|
|
query_img *= 255
|
|
query_img, query_pt = self.transformImg(query_img)
|
|
|
|
db_img *= 255
|
|
db_img, db_pt = self.transformImg(db_img)
|
|
|
|
return query_img, query_pt, db_img, db_pt, idx
|
|
|
|
def transformImg(self, img):
|
|
xs, ys = np.meshgrid(np.arange(self.pts_step,img.size()[1]-self.pts_step,self.pts_step), np.arange(self.pts_step,img.size()[2]-self.pts_step,self.pts_step))
|
|
xs=xs.reshape(-1,1)
|
|
ys = ys.reshape(-1,1)
|
|
pts = np.hstack((xs,ys))
|
|
img = img.permute(1,2,0).detach().numpy()
|
|
transformed_imgs=self.group_transformer.transform(img,pts)
|
|
data_img, data_pt = self.group_transformer.postprocess_transformed_imgs(transformed_imgs)
|
|
return data_img, data_pt
|
|
|
|
@staticmethod
|
|
def image_loader(path):
|
|
try:
|
|
return Image.open(path)
|
|
# return imread(path)
|
|
except UnidentifiedImageError:
|
|
print(f'Image {path} could not be loaded')
|
|
return Image.new('RGB', (224, 224))
|
|
|
|
def __len__(self):
|
|
|
|
return len(self.samples)
|
|
|
|
def shuffle(self, ):
|
|
|
|
"""
|
|
generate unique class_id
|
|
"""
|
|
print("\n Shuffle Dataset")
|
|
|
|
pair_pool = copy.deepcopy(self.pairs)
|
|
#shuffle
|
|
random.shuffle(pair_pool)
|
|
|
|
pairs_epoch = set()
|
|
label_batch = set()
|
|
|
|
current_batch = []
|
|
batches = []
|
|
|
|
# progressbar
|
|
pbar = tqdm()
|
|
|
|
while True:
|
|
pbar.update()
|
|
if len(pair_pool) > 0:
|
|
pair = pair_pool.pop(0)
|
|
|
|
label, _, _ = pair
|
|
|
|
if label not in label_batch and pair not in pairs_epoch:
|
|
|
|
label_batch.add(label)
|
|
current_batch.append(pair)
|
|
pairs_epoch.add(pair)
|
|
|
|
break_counter = 0
|
|
|
|
else:
|
|
if pair not in pairs_epoch:
|
|
pair_pool.append(pair)
|
|
|
|
break_counter += 1
|
|
|
|
if break_counter >= 5000:
|
|
break
|
|
|
|
else:
|
|
break
|
|
|
|
if len(current_batch) >= self.shuffle_batch_size:
|
|
batches.extend(current_batch)
|
|
label_batch = set()
|
|
current_batch = []
|
|
|
|
pbar.close()
|
|
|
|
time.sleep(0.3)
|
|
|
|
self.samples = batches
|
|
|
|
print("Original Length: {} - Length after Shuffle: {}".format(len(self.pairs), len(self.samples)))
|
|
print("Break Counter:", break_counter)
|
|
print("Pairs left out of last batch to avoid creating noise:", len(self.pairs) - len(self.samples))
|
|
# print("First Element ID: {} - Last Element ID: {}".format(self.samples[0][0], self.samples[-1][0]))
|
|
|
|
class WorldDatasetTrainVanilia(Dataset):
|
|
def __init__(self,
|
|
data_dir,
|
|
query_txt,
|
|
transforms_query=default_transform,
|
|
transforms_db=default_transform,
|
|
shuffle_batch_size=64):
|
|
super().__init__()
|
|
|
|
self.pairs = []
|
|
self.data = get_data(query_txt)
|
|
|
|
for idx in self.data.items():
|
|
query_img_path = os.path.join(data_dir, idx[1][0])
|
|
label = eval(idx[1][1])
|
|
db_image_path = os.path.join(data_dir, idx[1][2])
|
|
self.pairs.append((label, query_img_path, db_image_path))
|
|
|
|
self.transforms_query = transforms_query
|
|
self.transforms_db = transforms_db
|
|
self.shuffle_batch_size = shuffle_batch_size
|
|
|
|
self.samples = copy.deepcopy(self.pairs)
|
|
|
|
|
|
def __getitem__(self, index):
|
|
|
|
idx, query_img_path, db_img_path = self.samples[index]
|
|
# query
|
|
query_img = self.image_loader(query_img_path)
|
|
# db
|
|
db_img = self.image_loader(db_img_path)
|
|
# image transforms
|
|
if self.transforms_query is not None:
|
|
query_img = self.transforms_query(query_img)
|
|
|
|
if self.transforms_db is not None:
|
|
db_img = self.transforms_db(db_img)
|
|
|
|
return query_img, db_img, idx
|
|
|
|
@staticmethod
|
|
def image_loader(path):
|
|
try:
|
|
return Image.open(path)
|
|
# return imread(path)
|
|
except UnidentifiedImageError:
|
|
print(f'Image {path} could not be loaded')
|
|
return Image.new('RGB', (224, 224))
|
|
|
|
def __len__(self):
|
|
|
|
return len(self.samples)
|
|
|
|
def shuffle(self, ):
|
|
|
|
"""
|
|
generate unique class_id
|
|
"""
|
|
print("\n Shuffle Dataset")
|
|
|
|
pair_pool = copy.deepcopy(self.pairs)
|
|
#shuffle
|
|
random.shuffle(pair_pool)
|
|
|
|
pairs_epoch = set()
|
|
label_batch = set()
|
|
|
|
current_batch = []
|
|
batches = []
|
|
|
|
# progressbar
|
|
pbar = tqdm()
|
|
|
|
while True:
|
|
pbar.update()
|
|
if len(pair_pool) > 0:
|
|
pair = pair_pool.pop(0)
|
|
|
|
label, _, _ = pair
|
|
|
|
if label not in label_batch and pair not in pairs_epoch:
|
|
|
|
label_batch.add(label)
|
|
current_batch.append(pair)
|
|
pairs_epoch.add(pair)
|
|
|
|
break_counter = 0
|
|
|
|
else:
|
|
if pair not in pairs_epoch:
|
|
pair_pool.append(pair)
|
|
|
|
break_counter += 1
|
|
|
|
if break_counter >= 5000:
|
|
break
|
|
|
|
else:
|
|
break
|
|
|
|
if len(current_batch) >= self.shuffle_batch_size:
|
|
batches.extend(current_batch)
|
|
label_batch = set()
|
|
current_batch = []
|
|
|
|
pbar.close()
|
|
|
|
time.sleep(0.3)
|
|
|
|
self.samples = batches
|
|
|
|
print("Original Length: {} - Length after Shuffle: {}".format(len(self.pairs), len(self.samples)))
|
|
print("Break Counter:", break_counter)
|
|
print("Pairs left out of last batch to avoid creating noise:", len(self.pairs) - len(self.samples))
|
|
|
|
class WorldDatasetEvalGroup(Dataset):
|
|
def __init__(self,
|
|
data_dir,
|
|
name,
|
|
mode,
|
|
height_mode=None,
|
|
transforms=default_transform
|
|
):
|
|
super().__init__()
|
|
|
|
self.transforms = transforms
|
|
|
|
self.group_transformer = TransformerCV(transform_config)
|
|
self.pts_step = 5
|
|
|
|
self.data_dir = data_dir
|
|
self.name = name
|
|
|
|
pos_json_path = os.path.join(self.data_dir, self.name, 'positive.json')
|
|
positive = json.load(open(pos_json_path))
|
|
|
|
self.samples = []
|
|
if mode == 'query':
|
|
if os.path.exists(os.path.join(data_dir, name,'query', height_mode, 'footage')):
|
|
temp_path = os.path.join(data_dir, name,'query', height_mode, 'footage')
|
|
temp = sorted(glob.glob(f'{temp_path}/{"*.jpeg"}'))
|
|
if len(temp) != len(positive.keys()):
|
|
filter_temp = [image for image in temp if image.split('/')[-1].split('.')[0].split('_')[-1] in positive.keys()]
|
|
self.samples.extend(filter_temp)
|
|
else:
|
|
self.samples.extend(temp)
|
|
|
|
if mode == 'DB':
|
|
temp_path = os.path.join(data_dir, name, 'DB', 'img')
|
|
temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
|
|
self.samples.extend(temp)
|
|
|
|
|
|
|
|
|
|
def __getitem__(self, index):
|
|
|
|
img_path = self.samples[index]
|
|
# query
|
|
img = self.image_loader(img_path)
|
|
|
|
if self.transforms is not None:
|
|
img = self.transforms(img)
|
|
|
|
img *= 255
|
|
img, pt = self.transformImg(img)
|
|
|
|
return img, pt
|
|
|
|
def transformImg(self, img):
|
|
|
|
xs, ys = np.meshgrid(np.arange(self.pts_step,img.size()[1]-self.pts_step,self.pts_step), np.arange(self.pts_step,img.size()[2]-self.pts_step,self.pts_step))
|
|
xs=xs.reshape(-1,1)
|
|
ys = ys.reshape(-1,1)
|
|
pts = np.hstack((xs,ys))
|
|
img = img.permute(1,2,0).detach().numpy()
|
|
transformed_imgs=self.group_transformer.transform(img,pts)
|
|
data_img, data_pt = self.group_transformer.postprocess_transformed_imgs(transformed_imgs)
|
|
return data_img, data_pt
|
|
|
|
def get_gt(self,):
|
|
|
|
|
|
pos_json_path = os.path.join(self.data_dir, self.name, 'positive.json')
|
|
semi_pos_json_path = os.path.join(self.data_dir, self.name, 'semi_positive.json')
|
|
positive = json.load(open(pos_json_path))
|
|
semi_positive = json.load(open(semi_pos_json_path))
|
|
|
|
pos_gt = []
|
|
for key in positive.keys():
|
|
value = positive[key]
|
|
|
|
temp_index = []
|
|
# pos
|
|
for one_value in value:
|
|
temp_path_dir = os.path.join(self.data_dir, self.name, 'DB', 'img')
|
|
temp_path = temp_path_dir + '/' + one_value
|
|
one_index = self.samples.index(temp_path)
|
|
temp_index.append(one_index)
|
|
# semi-pos
|
|
try:
|
|
semi_value = semi_positive[key]
|
|
for one_value in semi_value:
|
|
temp_path_dir = os.path.join(self.data_dir, self.name, 'DB', 'img')
|
|
temp_path = temp_path_dir + '/' + one_value
|
|
one_index = self.samples.index(temp_path)
|
|
temp_index.append(one_index)
|
|
except:
|
|
pos_gt.append([key, temp_index])
|
|
continue
|
|
|
|
pos_gt.append([key, temp_index])
|
|
|
|
return pos_gt
|
|
|
|
|
|
def getitem(self, index):
|
|
|
|
return self.samples[index]
|
|
|
|
|
|
|
|
@staticmethod
|
|
def image_loader(path):
|
|
try:
|
|
return Image.open(path)
|
|
# return imread(path)
|
|
except UnidentifiedImageError:
|
|
print(f'Image {path} could not be loaded')
|
|
return Image.new('RGB', (224, 224))
|
|
|
|
def __len__(self):
|
|
|
|
return len(self.samples)
|
|
|
|
class WorldDatasetEvalVanilia(Dataset):
|
|
def __init__(self,
|
|
data_dir,
|
|
name,
|
|
mode,
|
|
height_mode=None,
|
|
transforms=default_transform
|
|
):
|
|
super().__init__()
|
|
|
|
self.transforms = transforms
|
|
|
|
self.data_dir = data_dir
|
|
self.name = name
|
|
|
|
pos_json_path = os.path.join(self.data_dir, self.name, 'positive.json')
|
|
positive = json.load(open(pos_json_path))
|
|
|
|
self.samples = []
|
|
if mode == 'query':
|
|
|
|
if os.path.exists(os.path.join(data_dir, name,'query', height_mode, 'footage')): #query
|
|
temp_path = os.path.join(data_dir, name,'query', height_mode, 'footage')
|
|
temp = sorted(glob.glob(f'{temp_path}/{"*.jpeg"}'))
|
|
if len(temp) != len(positive.keys()):
|
|
filter_temp = [image for image in temp if image.split('/')[-1].split('.')[0].split('_')[-1] in positive.keys()]
|
|
self.samples.extend(filter_temp)
|
|
else:
|
|
self.samples.extend(temp)
|
|
|
|
if mode == 'DB':
|
|
temp_path = os.path.join(data_dir, name, 'DB', 'img')
|
|
temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
|
|
self.samples.extend(temp)
|
|
|
|
|
|
|
|
|
|
def __getitem__(self, index):
|
|
|
|
img_path = self.samples[index]
|
|
# query
|
|
img = self.image_loader(img_path)
|
|
|
|
if self.transforms is not None:
|
|
img = self.transforms(img)
|
|
|
|
return img
|
|
|
|
|
|
def get_gt(self,):
|
|
|
|
|
|
pos_json_path = os.path.join(self.data_dir, self.name, 'positive.json')
|
|
semi_pos_json_path = os.path.join(self.data_dir, self.name, 'semi_positive.json')
|
|
positive = json.load(open(pos_json_path))
|
|
semi_positive = json.load(open(semi_pos_json_path))
|
|
|
|
pos_gt = []
|
|
for key in positive.keys():
|
|
value = positive[key]
|
|
|
|
temp_index = []
|
|
# pos
|
|
for one_value in value:
|
|
temp_path_dir = os.path.join(self.data_dir, self.name, 'DB', 'img')
|
|
temp_path = temp_path_dir + '/' + one_value
|
|
one_index = self.samples.index(temp_path)
|
|
temp_index.append(one_index)
|
|
try:
|
|
semi_value = semi_positive[key]
|
|
# semi-pos
|
|
for one_value in semi_value:
|
|
temp_path_dir = os.path.join(self.data_dir, self.name, 'DB', 'img')
|
|
temp_path = temp_path_dir + '/' + one_value
|
|
one_index = self.samples.index(temp_path)
|
|
temp_index.append(one_index)
|
|
except:
|
|
pos_gt.append([key, temp_index])
|
|
continue
|
|
|
|
pos_gt.append([key, temp_index])
|
|
|
|
return pos_gt
|
|
|
|
|
|
def getitem(self, index):
|
|
|
|
return self.samples[index]
|
|
|
|
|
|
|
|
@staticmethod
|
|
def image_loader(path):
|
|
try:
|
|
return Image.open(path)
|
|
# return imread(path)
|
|
except UnidentifiedImageError:
|
|
print(f'Image {path} could not be loaded')
|
|
return Image.new('RGB', (224, 224))
|
|
|
|
def __len__(self):
|
|
|
|
return len(self.samples)
|
|
|
|
class AerialDatasetEvalVanilia(Dataset):
|
|
def __init__(self,
|
|
data_dir,
|
|
mode,
|
|
angle=0,
|
|
transforms=default_transform
|
|
):
|
|
super().__init__()
|
|
|
|
self.samples = []
|
|
if mode == 'query':
|
|
temp_path = os.path.join(data_dir, 'query_images')
|
|
temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
|
|
self.samples.extend(temp)
|
|
self.angle = angle
|
|
|
|
if mode == 'DB':
|
|
temp_path = os.path.join(data_dir, 'reference_images')
|
|
temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
|
|
self.samples.extend(temp)
|
|
self.angle = angle
|
|
|
|
self.transforms = transforms
|
|
self.data_dir = data_dir
|
|
self.mode = mode
|
|
|
|
|
|
def __getitem__(self, index):
|
|
|
|
img_path = self.samples[index]
|
|
# query
|
|
img = self.image_loader(img_path, self.mode, self.angle)
|
|
|
|
if self.transforms is not None:
|
|
img = self.transforms(img)
|
|
|
|
return img
|
|
|
|
|
|
def get_gt(self,):
|
|
|
|
columns_to_use_by_index = [1, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20,]
|
|
pos_cvs_path = os.path.join(self.data_dir, 'gt_matches.csv')
|
|
df = pd.read_csv(pos_cvs_path, usecols=columns_to_use_by_index)
|
|
|
|
pos_gt = []
|
|
for i in range(len(df)):
|
|
for j in range(df.shape[1]):
|
|
if j == 0:
|
|
key = df.iloc[i, j]
|
|
temp_index = []
|
|
else:
|
|
value = df.iloc[i, j]
|
|
temp_index.append(value)
|
|
|
|
pos_gt.append([key, temp_index])
|
|
return pos_gt
|
|
|
|
def get_gt_npy(self,):
|
|
|
|
data_path = os.path.join(self.data_dir, 'vpair_gt.npy')
|
|
data = np.load(data_path, allow_pickle=True)
|
|
pos_gt = []
|
|
for i in range(data.shape[0]):
|
|
key = data[i, 0]
|
|
temp_index = []
|
|
temp_value = data[i, 1]
|
|
for j in temp_value:
|
|
temp_index.append(j)
|
|
|
|
pos_gt.append([key, temp_index])
|
|
return pos_gt
|
|
|
|
|
|
def getitem(self, index):
|
|
|
|
return self.samples[index]
|
|
|
|
|
|
|
|
@staticmethod
|
|
def image_loader(path, mode, angle):
|
|
try:
|
|
if mode == 'query':
|
|
img = Image.open(path)
|
|
if angle == 0:
|
|
return img
|
|
rotated_image = img.rotate(angle,expand=True)
|
|
return rotated_image
|
|
else:
|
|
return Image.open(path)
|
|
# Image.open(path)
|
|
# return imread(path)
|
|
except UnidentifiedImageError:
|
|
print(f'Image {path} could not be loaded')
|
|
return Image.new('RGB', (224, 224))
|
|
|
|
def __len__(self):
|
|
|
|
return len(self.samples)
|
|
|
|
class AerialDatasetEvalGroup(Dataset):
|
|
def __init__(self,
|
|
data_dir,
|
|
mode,
|
|
angle=0,
|
|
transforms=default_transform
|
|
):
|
|
super().__init__()
|
|
|
|
self.samples = []
|
|
if mode == 'query':
|
|
temp_path = os.path.join(data_dir, 'query_images')
|
|
temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
|
|
self.samples.extend(temp)
|
|
self.angle = angle
|
|
|
|
if mode == 'DB':
|
|
temp_path = os.path.join(data_dir, 'reference_images')
|
|
temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
|
|
self.samples.extend(temp)
|
|
self.angle = 0
|
|
|
|
self.transforms = transforms
|
|
self.mode = mode
|
|
|
|
self.group_transformer = TransformerCV(transform_config)
|
|
self.pts_step = 5
|
|
|
|
|
|
self.data_dir = data_dir
|
|
|
|
|
|
def __getitem__(self, index):
|
|
|
|
img_path = self.samples[index]
|
|
# query
|
|
img = self.image_loader(img_path, self.mode, self.angle)
|
|
|
|
if self.transforms is not None:
|
|
img = self.transforms(img)
|
|
|
|
# group
|
|
img *= 255
|
|
img, pt = self.transformImg(img)
|
|
|
|
return img, pt
|
|
|
|
def transformImg(self, img):
|
|
|
|
xs, ys = np.meshgrid(np.arange(self.pts_step,img.size()[1]-self.pts_step,self.pts_step), np.arange(self.pts_step,img.size()[2]-self.pts_step,self.pts_step))
|
|
xs=xs.reshape(-1,1)
|
|
ys = ys.reshape(-1,1)
|
|
pts = np.hstack((xs,ys))
|
|
img = img.permute(1,2,0).detach().numpy()
|
|
transformed_imgs=self.group_transformer.transform(img,pts)
|
|
data_img, data_pt = self.group_transformer.postprocess_transformed_imgs(transformed_imgs)
|
|
return data_img, data_pt
|
|
|
|
def get_gt(self,):
|
|
|
|
columns_to_use_by_index = [1, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20,]
|
|
pos_cvs_path = os.path.join(self.data_dir, 'gt_matches.csv')
|
|
df = pd.read_csv(pos_cvs_path, usecols=columns_to_use_by_index)
|
|
|
|
pos_gt = []
|
|
for i in range(len(df)):
|
|
for j in range(df.shape[1]):
|
|
if j == 0:
|
|
key = df.iloc[i, j]
|
|
temp_index = []
|
|
else:
|
|
value = df.iloc[i, j]
|
|
temp_index.append(value)
|
|
|
|
pos_gt.append([key, temp_index])
|
|
return pos_gt
|
|
|
|
def get_gt_npy(self,):
|
|
|
|
data_path = os.path.join(self.data_dir, 'vpair_gt.npy')
|
|
data = np.load(data_path, allow_pickle=True)
|
|
pos_gt = []
|
|
for i in range(data.shape[0]):
|
|
key = data[i, 0]
|
|
temp_index = []
|
|
temp_value = data[i, 1]
|
|
for j in temp_value:
|
|
temp_index.append(j)
|
|
|
|
pos_gt.append([key, temp_index])
|
|
return pos_gt
|
|
|
|
|
|
def getitem(self, index):
|
|
|
|
return self.samples[index]
|
|
|
|
|
|
|
|
@staticmethod
|
|
def image_loader(path, mode, angle):
|
|
try:
|
|
if mode == 'query':
|
|
img = Image.open(path)
|
|
rotated_image = img.rotate(angle,expand=True)
|
|
return rotated_image
|
|
else:
|
|
return Image.open(path)
|
|
except UnidentifiedImageError:
|
|
print(f'Image {path} could not be loaded')
|
|
return Image.new('RGB', (224, 224))
|
|
|
|
def __len__(self):
|
|
|
|
return len(self.samples)
|
|
|
|
|
|
class DenseUAVDatasetEvalVanilia(Dataset):
|
|
def __init__(self,
|
|
txt,
|
|
mode,
|
|
gt_txt,
|
|
transforms=default_transform
|
|
):
|
|
super().__init__()
|
|
|
|
self.samples = []
|
|
if mode == 'query':
|
|
with open(txt, 'r') as f:
|
|
for i in f:
|
|
self.samples.append(i.strip())
|
|
|
|
|
|
if mode == 'DB':
|
|
with open(txt, 'r') as f:
|
|
for i in f:
|
|
self.samples.append(i.strip())
|
|
|
|
self.transforms = transforms
|
|
self.mode = mode
|
|
self.gt_txt = gt_txt
|
|
|
|
|
|
def __getitem__(self, index):
|
|
|
|
img_path = self.samples[index]
|
|
# query
|
|
img = self.image_loader(img_path, self.mode)
|
|
|
|
if self.transforms is not None:
|
|
img = self.transforms(img)
|
|
|
|
return img
|
|
|
|
|
|
def get_gt(self,):
|
|
|
|
pos_gt = []
|
|
with open(self.gt_txt, 'r') as f_gt:
|
|
for info in f_gt:
|
|
key, values = eval(info.split(' ')[0]), info.strip().split(' ')[1:]
|
|
temp_value = []
|
|
for value in values:
|
|
temp_value.append(eval(value))
|
|
pos_gt.append([key, temp_value])
|
|
|
|
return pos_gt
|
|
|
|
|
|
def getitem(self, index):
|
|
|
|
return self.samples[index]
|
|
|
|
|
|
|
|
@staticmethod
|
|
def image_loader(path, mode):
|
|
try:
|
|
if mode == 'query':
|
|
img = Image.open(path)
|
|
rotated_image = img.rotate(0,expand=True)
|
|
return rotated_image
|
|
else:
|
|
return Image.open(path)
|
|
# Image.open(path)
|
|
# return imread(path)
|
|
except UnidentifiedImageError:
|
|
print(f'Image {path} could not be loaded')
|
|
return Image.new('RGB', (224, 224))
|
|
|
|
def __len__(self):
|
|
|
|
return len(self.samples)
|
|
|
|
class DenseUAVDatasetEvalGroup(Dataset):
|
|
def __init__(self,
|
|
txt,
|
|
mode,
|
|
gt_txt,
|
|
transforms=default_transform
|
|
):
|
|
super().__init__()
|
|
|
|
self.samples = []
|
|
if mode == 'query':
|
|
with open(txt, 'r') as f:
|
|
for i in f:
|
|
self.samples.append(i.strip())
|
|
|
|
|
|
if mode == 'DB':
|
|
with open(txt, 'r') as f:
|
|
for i in f:
|
|
self.samples.append(i.strip())
|
|
|
|
self.transforms = transforms
|
|
self.mode = mode
|
|
|
|
self.group_transformer = TransformerCV(transform_config)
|
|
self.pts_step = 5
|
|
self.gt_txt = gt_txt
|
|
|
|
|
|
def __getitem__(self, index):
|
|
|
|
img_path = self.samples[index]
|
|
# query
|
|
img = self.image_loader(img_path, self.mode)
|
|
|
|
if self.transforms is not None:
|
|
img = self.transforms(img)
|
|
|
|
# group
|
|
img *= 255
|
|
img, pt = self.transformImg(img)
|
|
|
|
return img, pt
|
|
|
|
def transformImg(self, img):
|
|
|
|
xs, ys = np.meshgrid(np.arange(self.pts_step,img.size()[1]-self.pts_step,self.pts_step), np.arange(self.pts_step,img.size()[2]-self.pts_step,self.pts_step))
|
|
xs=xs.reshape(-1,1)
|
|
ys = ys.reshape(-1,1)
|
|
pts = np.hstack((xs,ys))
|
|
img = img.permute(1,2,0).detach().numpy()
|
|
transformed_imgs=self.group_transformer.transform(img,pts)
|
|
data_img, data_pt = self.group_transformer.postprocess_transformed_imgs(transformed_imgs)
|
|
return data_img, data_pt
|
|
|
|
def get_gt(self,):
|
|
|
|
pos_gt = []
|
|
with open(self.gt_txt, 'r') as f_gt:
|
|
for info in f_gt:
|
|
key, values = eval(info.split(' ')[0]), info.strip().split(' ')[1:]
|
|
temp_value = []
|
|
for value in values:
|
|
temp_value.append(eval(value))
|
|
pos_gt.append([key, temp_value])
|
|
|
|
return pos_gt
|
|
|
|
def getitem(self, index):
|
|
|
|
return self.samples[index]
|
|
|
|
|
|
|
|
@staticmethod
|
|
def image_loader(path, mode):
|
|
try:
|
|
if mode == 'query':
|
|
img = Image.open(path)
|
|
rotated_image = img.rotate(270,expand=True)
|
|
return rotated_image
|
|
else:
|
|
return Image.open(path)
|
|
except UnidentifiedImageError:
|
|
print(f'Image {path} could not be loaded')
|
|
return Image.new('RGB', (224, 224))
|
|
|
|
def __len__(self):
|
|
|
|
return len(self.samples)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# 测试代码
|
|
|
|
# data_dir = "/media/guan/新加卷/EdgeBing/WorldLoc"
|
|
# query_txt = "/media/guan/新加卷/EdgeBing/WorldLoc/Index/train_query.txt"
|
|
|
|
# train_dataset = WorldDatasetTrain(data_dir, query_txt)
|
|
|
|
# train_dataloader = DataLoader(train_dataset,
|
|
# batch_size=64,
|
|
# num_workers=0,
|
|
# shuffle=False,
|
|
# pin_memory=True)
|
|
|
|
|
|
# train_dataloader.dataset.shuffle()
|
|
|
|
# for query, query_pt, reference, reference_pt, idx in tqdm(train_dataloader, total=len(train_dataloader)):
|
|
|
|
# print(1) |