started working on new retrieval function

parent 5f165143
......@@ -15,7 +15,7 @@ RANDOM_CROP_SITE_FINETUNING = 224
CENTER_CROP_SIZE_KEYPOINT = 900
RANDOM_CROP_SITE_KEYPOINT = 720
BATCH_PRINT_FREQ = 20
MAX_PCA_FEATURES = 500000
class ExperimentManager:
"""
......@@ -129,10 +129,10 @@ class ExperimentManager:
for stage in self.dataset.keys():
assert self.dataset[stage].is_dir(), f"Dataset {self.dataset[stage]} for {stage} not found!"
if stage == "retrieval":
assert self.dataset[stage].joinpath(Path("index")).is_dir(), f"Dataset for retrieval must contain " \
f"index subdirectory"
assert self.dataset[stage].joinpath(Path("query")).is_dir(), f"Dataset for retrieval must contain " \
f"query subdirectory"
assert self.dataset[stage].joinpath(Path("data")).is_dir(), f"Dataset for retrieval must contain " \
f"data subdirectory"
assert self.dataset[stage].joinpath(Path("query.json")).is_file(), f"Dataset for retrieval must " \
f"contain query file"
print("All dataset directories found!")
# create load and save directories for the created/ used models. Automatically infer paths if multiple stages
......@@ -144,7 +144,8 @@ class ExperimentManager:
if "keypoints" in self.stages:
self.train_model("keypoints")
if "retrieval" in self.stages:
self.perform_retrieval()
#self.perform_retrieval()
self.perform_retrieval_2()
def get_load_and_save_paths(self, load_from):
"""
......@@ -343,6 +344,57 @@ class ExperimentManager:
json.dump(log, json_file, ensure_ascii=False)
print(f"Completed {stage} training!")
def perform_retrieval_2(self):
# load model in retrieval mode
model = Delf(None, "retrieval", self.load_paths["retrieval"], target_layer=self.target_layer,
use_l2_normalization=self.use_l2_normalization).to(self.device)
model.eval()
# get path to all images
images = self.dataset["retrieval"].joinpath(Path("data"))
# get path to query list
query = self.dataset["retrieval"].joinpath(Path("query.json"))
# create path to save results
stage_path = self.experiment_path.joinpath(Path("retrieval"))
with torch.no_grad():
accumulated_features = None
image_info = {}
data_loader = get_path_data_loader(images, num_workers=self.num_workers["retrieval"])
for ind, batch in enumerate(data_loader):
# get all image info from loader
image_input, label, path = batch
image_input = image_input.to(self.device)
filename = Path(path[0]).stem
# get features and rf info from model
features, rf_centers = model.extract_features(image_input)
# Delf says normalize before pca
features = torch.nn.functional.normalize(features, p=2, dim=1)
# accumulate features for pca_calculations
if accumulated_features is None:
accumulated_features = features
else:
accumulated_features = torch.cat((accumulated_features, features), dim=0)
image_info[filename] = [features, rf_centers]
accumulated_features = accumulated_features.numpy()
# if we have too many features for pca we choose a random subset
if len(accumulated_features) > MAX_PCA_FEATURES:
accumulated_features = np.random.permutation(accumulated_features)[:MAX_PCA_FEATURES]
# calculate and evaluate the pca matrix
pca = calculate_pca(accumulated_features, stage_path.joinpath(
f"{self.load_paths['retrieval'].stem}_{self.dataset['retrieval'].stem}.pca"), self.pca_log)
for filename in image_info:
features, rf_centers = image_info[filename]
rf_centers = rf_centers.numpy()
features = pca.transform(features.numpy())
if self.use_retrieval_normalization:
features = normalize(features, norm='l2', axis=1)
image_info[filename] = [features, rf_centers]
"""
read from query file
multithread matching mit workern dafür dann kdtree in single thread stellen
ausgabe als json file pro query mit tripeln (filename, matching score, ransac score)
"""
def perform_retrieval(self):
# load model in retrieval mode
model = Delf(None, "retrieval", self.load_paths["retrieval"], target_layer=self.target_layer,
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment