started working on new retrieval function

parent 5f165143
......@@ -15,7 +15,7 @@ RANDOM_CROP_SITE_FINETUNING = 224
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:
if "retrieval" in self.stages:
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,
# 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 =
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
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
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment