Integrated the previous changes with RetrievalResults class. Removed keys. Changed signature of EmbeddingMetrics.

**CHANGELOG**

* removed keys: `IS_QUERY_KEY`, `IS_GALLERY_KEY`, `CATEGORIES_KEY`, `PATHS_KEY`, `X1_KEY`, `X2_KEY`, `Y1_KEY`, `Y2_KEY`, `SEQUENCE_KEY`. Categories and sequences are passed through `extra_data` instead. The rest is incapsulated in Dataset.
* Removed `IMetricDDP`, `EmbeddingMetricsDDP`. Reason: having `EmbeddingMetrics` is enough, because we do accumulator sync there anyway.
* Changed signatures of `EmbeddingMetrics`: keys replaces by providing dataset, removed `.sync()` and `.visualisation()` methods and so on.
* Updated `.md` examples and `.rst` docs

Minor:
* removed: `calc_distance_matrix`, `validate_dataset`  -- this logic happens in `RetrievalResults`, also removed `find_first_occurrences` -- we have `unique_by_ids` instead.
* removed `DummyDataset` in tests (used `EmbeddingsQueryGalleryLabeledDataset` instead).

Author: AlekseySh

Makefile

@@ -119,6 +119,11 @@ pip_install_actual_oml:
 .PHONY: clean
 clean:
 	find . -type d -name "__pycache__" -exec rm -r {} +
+	find . -type d -name "lightning_logs" -exec rm -r {} +
+	find . -type d -name "ml-runs" -exec rm -r {} +
+	find . -type d -name "logs" -exec rm -r {} +
+	find . -type d -name ".ipynb_checkpoints" -exec rm -r {} +
 	find . -type f -name "*.log" -exec rm {} +
-	find . -type f -name "*.predictions.json" -exec rm {} +
+	find . -type f -name "*predictions.json" -exec rm {} +
 	rm -rf docs/build
+	rm -rf outputs/

README.md

@@ -294,7 +294,7 @@ docker pull omlteam/oml:cpu
 import torch
 from tqdm import tqdm
 
-from oml.datasets.base import DatasetWithLabels
+from oml.datasets import ImageLabeledDataset
 from oml.losses.triplet import TripletLossWithMiner
 from oml.miners.inbatch_all_tri import AllTripletsMiner
 from oml.models import ViTExtractor
@@ -306,7 +306,7 @@ df_train, _ = download_mock_dataset(global_paths=True)
 extractor = ViTExtractor("vits16_dino", arch="vits16", normalise_features=False).train()
 optimizer = torch.optim.SGD(extractor.parameters(), lr=1e-6)
 
-train_dataset = DatasetWithLabels(df_train)
+train_dataset = ImageLabeledDataset(df_train)
 criterion = TripletLossWithMiner(margin=0.1, miner=AllTripletsMiner(), need_logs=True)
 sampler = BalanceSampler(train_dataset.get_labels(), n_labels=2, n_instances=2)
 train_loader = torch.utils.data.DataLoader(train_dataset, batch_sampler=sampler)
@@ -333,39 +333,28 @@ for batch in tqdm(train_loader):
 
 [comment]:vanilla-validation-start
 ```python
-import torch
-from tqdm import tqdm
 
-from oml.datasets.base import DatasetQueryGallery
-from oml.metrics.embeddings import EmbeddingMetrics
+from oml.datasets import ImageQueryGalleryLabeledDataset
+from oml.inference import inference
+from oml.metrics import calc_retrieval_metrics_rr
 from oml.models import ViTExtractor
+from oml.retrieval import RetrievalResults
 from oml.utils.download_mock_dataset import download_mock_dataset
+from oml.registry.transforms import get_transforms_for_pretrained
 
-_, df_val = download_mock_dataset(global_paths=True)
-
-extractor = ViTExtractor("vits16_dino", arch="vits16", normalise_features=False).eval()
-
-val_dataset = DatasetQueryGallery(df_val)
-
-val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=4)
-calculator = EmbeddingMetrics(extra_keys=("paths",))
-calculator.setup(num_samples=len(val_dataset))
+extractor = ViTExtractor.from_pretrained("vits16_dino")
+transform, _ = get_transforms_for_pretrained("vits16_dino")
 
-with torch.no_grad():
-    for batch in tqdm(val_loader):
-        batch["embeddings"] = extractor(batch["input_tensors"])
-        calculator.update_data(batch)
+_, df_val = download_mock_dataset(global_paths=True)
+dataset = ImageQueryGalleryLabeledDataset(df_val, transform=transform)
 
-metrics = calculator.compute_metrics()
+embeddings = inference(extractor, dataset, batch_size=4)
 
-# Logging
-print(calculator.metrics)  # metrics
-print(calculator.metrics_unreduced)  # metrics without averaging over queries
+rr = RetrievalResults.compute_from_embeddings(embeddings, dataset, n_items_to_retrieve=5)
+metrics = calc_retrieval_metrics_rr(rr, map_top_k=(3, 5), precision_top_k=(5,), cmc_top_k=(3,))
 
-# Visualisation
-calculator.get_plot_for_queries(query_ids=[0, 2], n_instances=5)  # draw predictions on predefined queries
-calculator.get_plot_for_worst_queries(metric_name="OVERALL/map/5", n_queries=2, n_instances=5)  # draw mistakes
-calculator.visualize()  # draw mistakes for all the available metrics
+print(rr, "\n", metrics)
+rr.visualize(query_ids=[2, 1], dataset=dataset).show()
 
 ```
 [comment]:vanilla-validation-end
@@ -380,9 +369,10 @@ calculator.visualize()  # draw mistakes for all the available metrics
 [comment]:lightning-start
 ```python
 import pytorch_lightning as pl
-import torch
+from torch.utils.data import DataLoader
+from torch.optim import SGD
 
-from oml.datasets.base import DatasetQueryGallery, DatasetWithLabels
+from oml.datasets import ImageLabeledDataset, ImageQueryGalleryLabeledDataset
 from oml.lightning.modules.extractor import ExtractorModule
 from oml.lightning.callbacks.metric import MetricValCallback
 from oml.losses.triplet import TripletLossWithMiner
@@ -405,16 +395,16 @@ df_train, df_val = download_mock_dataset(global_paths=True)
 extractor = ViTExtractor("vits16_dino", arch="vits16", normalise_features=False)
 
 # train
-optimizer = torch.optim.SGD(extractor.parameters(), lr=1e-6)
-train_dataset = DatasetWithLabels(df_train)
+optimizer = SGD(extractor.parameters(), lr=1e-6)
+train_dataset = ImageLabeledDataset(df_train)
 criterion = TripletLossWithMiner(margin=0.1, miner=AllTripletsMiner())
 batch_sampler = BalanceSampler(train_dataset.get_labels(), n_labels=2, n_instances=3)
-train_loader = torch.utils.data.DataLoader(train_dataset, batch_sampler=batch_sampler)
+train_loader = DataLoader(train_dataset, batch_sampler=batch_sampler)
 
 # val
-val_dataset = DatasetQueryGallery(df_val)
-val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=4)
-metric_callback = MetricValCallback(metric=EmbeddingMetrics(extra_keys=[train_dataset.paths_key,]), log_images=True)
+val_dataset = ImageQueryGalleryLabeledDataset(df_val)
+val_loader = DataLoader(val_dataset, batch_size=4)
+metric_callback = MetricValCallback(metric=EmbeddingMetrics(dataset=val_dataset), log_images=True)
 
 # 1) Logging with Tensorboard
 logger = TensorBoardPipelineLogger(".")
@@ -450,14 +440,13 @@ trainer.fit(pl_model, train_dataloaders=train_loader, val_dataloaders=val_loader
 
 [comment]:usage-retrieval-start
 ```python
-import torch
-
 from oml.datasets import ImageQueryGalleryDataset
 from oml.inference import inference
 from oml.models import ViTExtractor
 from oml.registry.transforms import get_transforms_for_pretrained
 from oml.utils.download_mock_dataset import download_mock_dataset
-from oml.utils.misc_torch import pairwise_dist
+from oml.retrieval.retrieval_results import RetrievalResults
+
 
 _, df_test = download_mock_dataset(global_paths=True)
 del df_test["label"]  # we don't need gt labels for doing predictions
@@ -466,25 +455,14 @@ extractor = ViTExtractor.from_pretrained("vits16_dino")
 transform, _ = get_transforms_for_pretrained("vits16_dino")
 
 dataset = ImageQueryGalleryDataset(df_test, transform=transform)
-
 embeddings = inference(extractor, dataset, batch_size=4, num_workers=0)
-embeddings_query = embeddings[dataset.get_query_ids()]
-embeddings_gallery = embeddings[dataset.get_gallery_ids()]
 
-# Now we can explicitly build pairwise matrix of distances or save you RAM via using kNN
-use_knn = False
-top_k = 3
+retrieval_results = RetrievalResults.compute_from_embeddings(embeddings, dataset, n_items_to_retrieve=5)
 
-if use_knn:
-    from sklearn.neighbors import NearestNeighbors
-    knn = NearestNeighbors(algorithm="auto", p=2).fit(embeddings_query)
-    dists, ii_closest = knn.kneighbors(embeddings_gallery, n_neighbors=top_k, return_distance=True)
+retrieval_results.visualize(query_ids=[0, 1], dataset=dataset).show()
 
-else:
-    dist_mat = pairwise_dist(x1=embeddings_query, x2=embeddings_gallery, p=2)
-    dists, ii_closest = torch.topk(dist_mat, dim=1, k=top_k, largest=False)
+print(retrieval_results)  # you get the ids of retrieved items and the corresponding distances
 
-print(f"Top {top_k} items closest to queries are:\n {ii_closest}")
 ```
 [comment]:usage-retrieval-end
 </p>

docs/readme/examples_source/extractor/retrieval_usage.md

@@ -4,14 +4,13 @@
 
 [comment]:usage-retrieval-start
 ```python
-import torch
-
 from oml.datasets import ImageQueryGalleryDataset
 from oml.inference import inference
 from oml.models import ViTExtractor
 from oml.registry.transforms import get_transforms_for_pretrained
 from oml.utils.download_mock_dataset import download_mock_dataset
-from oml.utils.misc_torch import pairwise_dist
+from oml.retrieval.retrieval_results import RetrievalResults
+
 
 _, df_test = download_mock_dataset(global_paths=True)
 del df_test["label"]  # we don't need gt labels for doing predictions
@@ -20,25 +19,14 @@ extractor = ViTExtractor.from_pretrained("vits16_dino")
 transform, _ = get_transforms_for_pretrained("vits16_dino")
 
 dataset = ImageQueryGalleryDataset(df_test, transform=transform)
-
 embeddings = inference(extractor, dataset, batch_size=4, num_workers=0)
-embeddings_query = embeddings[dataset.get_query_ids()]
-embeddings_gallery = embeddings[dataset.get_gallery_ids()]
 
-# Now we can explicitly build pairwise matrix of distances or save you RAM via using kNN
-use_knn = False
-top_k = 3
+retrieval_results = RetrievalResults.compute_from_embeddings(embeddings, dataset, n_items_to_retrieve=5)
 
-if use_knn:
-    from sklearn.neighbors import NearestNeighbors
-    knn = NearestNeighbors(algorithm="auto", p=2).fit(embeddings_query)
-    dists, ii_closest = knn.kneighbors(embeddings_gallery, n_neighbors=top_k, return_distance=True)
+retrieval_results.visualize(query_ids=[0, 1], dataset=dataset).show()
 
-else:
-    dist_mat = pairwise_dist(x1=embeddings_query, x2=embeddings_gallery, p=2)
-    dists, ii_closest = torch.topk(dist_mat, dim=1, k=top_k, largest=False)
+print(retrieval_results)  # you get the ids of retrieved items and the corresponding distances
 
-print(f"Top {top_k} items closest to queries are:\n {ii_closest}")
 ```
 [comment]:usage-retrieval-end
 </p>

docs/readme/examples_source/extractor/train.md

@@ -7,7 +7,7 @@
 import torch
 from tqdm import tqdm
 
-from oml.datasets.base import DatasetWithLabels
+from oml.datasets import ImageLabeledDataset
 from oml.losses.triplet import TripletLossWithMiner
 from oml.miners.inbatch_all_tri import AllTripletsMiner
 from oml.models import ViTExtractor
@@ -19,7 +19,7 @@ df_train, _ = download_mock_dataset(global_paths=True)
 extractor = ViTExtractor("vits16_dino", arch="vits16", normalise_features=False).train()
 optimizer = torch.optim.SGD(extractor.parameters(), lr=1e-6)
 
-train_dataset = DatasetWithLabels(df_train)
+train_dataset = ImageLabeledDataset(df_train)
 criterion = TripletLossWithMiner(margin=0.1, miner=AllTripletsMiner(), need_logs=True)
 sampler = BalanceSampler(train_dataset.get_labels(), n_labels=2, n_instances=2)
 train_loader = torch.utils.data.DataLoader(train_dataset, batch_sampler=sampler)

docs/readme/examples_source/extractor/train_2loaders_val.md

@@ -5,9 +5,10 @@
 [comment]:lightning-2loaders-start
 ```python
 import pytorch_lightning as pl
-import torch
 
-from oml.datasets.base import DatasetQueryGallery
+from torch.utils.data import DataLoader
+
+from oml.datasets import ImageQueryGalleryLabeledDataset
 from oml.lightning.callbacks.metric import MetricValCallback
 from oml.lightning.modules.extractor import ExtractorModule
 from oml.metrics.embeddings import EmbeddingMetrics
@@ -20,24 +21,24 @@ _, df_val = download_mock_dataset(global_paths=True)
 extractor = ViTExtractor("vits16_dino", arch="vits16", normalise_features=False)
 
 # 1st validation dataset (big images)
-val_dataset_1 = DatasetQueryGallery(df_val, transform=get_normalisation_resize_torch(im_size=224))
-val_loader_1 = torch.utils.data.DataLoader(val_dataset_1, batch_size=4)
-metric_callback_1 = MetricValCallback(metric=EmbeddingMetrics(extra_keys=[val_dataset_1.paths_key,]),
+val_dataset_1 = ImageQueryGalleryLabeledDataset(df_val, transform=get_normalisation_resize_torch(im_size=224))
+val_loader_1 = DataLoader(val_dataset_1, batch_size=4)
+metric_callback_1 = MetricValCallback(metric=EmbeddingMetrics(dataset=val_dataset_1),
                                       log_images=True, loader_idx=0)
 
 # 2nd validation dataset (small images)
-val_dataset_2 = DatasetQueryGallery(df_val, transform=get_normalisation_resize_torch(im_size=48))
-val_loader_2 = torch.utils.data.DataLoader(val_dataset_2, batch_size=4)
-metric_callback_2 = MetricValCallback(metric=EmbeddingMetrics(extra_keys=[val_dataset_2.paths_key,]),
+val_dataset_2 = ImageQueryGalleryLabeledDataset(df_val, transform=get_normalisation_resize_torch(im_size=48))
+val_loader_2 = DataLoader(val_dataset_2, batch_size=4)
+metric_callback_2 = MetricValCallback(metric=EmbeddingMetrics(dataset=val_dataset_2),
                                       log_images=True, loader_idx=1)
 
 # run validation
 pl_model = ExtractorModule(extractor, None, None)
 trainer = pl.Trainer(max_epochs=3, callbacks=[metric_callback_1, metric_callback_2], num_sanity_val_steps=0)
 trainer.validate(pl_model, dataloaders=(val_loader_1, val_loader_2))
 
-print(metric_callback_1.metric.metrics)
-print(metric_callback_2.metric.metrics)
+print(metric_callback_1.metric.retrieval_results)
+print(metric_callback_2.metric.retrieval_results)
 ```
 [comment]:lightning-2loaders-end
 </p>

docs/readme/examples_source/extractor/train_val_pl.md

@@ -5,9 +5,10 @@
 [comment]:lightning-start
 ```python
 import pytorch_lightning as pl
-import torch
+from torch.utils.data import DataLoader
+from torch.optim import SGD
 
-from oml.datasets.base import DatasetQueryGallery, DatasetWithLabels
+from oml.datasets import ImageLabeledDataset, ImageQueryGalleryLabeledDataset
 from oml.lightning.modules.extractor import ExtractorModule
 from oml.lightning.callbacks.metric import MetricValCallback
 from oml.losses.triplet import TripletLossWithMiner
@@ -30,16 +31,16 @@ df_train, df_val = download_mock_dataset(global_paths=True)
 extractor = ViTExtractor("vits16_dino", arch="vits16", normalise_features=False)
 
 # train
-optimizer = torch.optim.SGD(extractor.parameters(), lr=1e-6)
-train_dataset = DatasetWithLabels(df_train)
+optimizer = SGD(extractor.parameters(), lr=1e-6)
+train_dataset = ImageLabeledDataset(df_train)
 criterion = TripletLossWithMiner(margin=0.1, miner=AllTripletsMiner())
 batch_sampler = BalanceSampler(train_dataset.get_labels(), n_labels=2, n_instances=3)
-train_loader = torch.utils.data.DataLoader(train_dataset, batch_sampler=batch_sampler)
+train_loader = DataLoader(train_dataset, batch_sampler=batch_sampler)
 
 # val
-val_dataset = DatasetQueryGallery(df_val)
-val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=4)
-metric_callback = MetricValCallback(metric=EmbeddingMetrics(extra_keys=[train_dataset.paths_key,]), log_images=True)
+val_dataset = ImageQueryGalleryLabeledDataset(df_val)
+val_loader = DataLoader(val_dataset, batch_size=4)
+metric_callback = MetricValCallback(metric=EmbeddingMetrics(dataset=val_dataset), log_images=True)
 
 # 1) Logging with Tensorboard
 logger = TensorBoardPipelineLogger(".")

docs/readme/examples_source/extractor/train_val_pl_ddp.md

@@ -6,13 +6,14 @@
 [comment]:lightning-ddp-start
 ```python
 import pytorch_lightning as pl
-import torch
+from torch.utils.data import DataLoader
+from torch.optim import SGD
 
-from oml.datasets.base import DatasetQueryGallery, DatasetWithLabels
+from oml.datasets import ImageQueryGalleryLabeledDataset, ImageLabeledDataset
 from oml.lightning.modules.extractor import ExtractorModuleDDP
 from oml.lightning.callbacks.metric import MetricValCallbackDDP
 from oml.losses.triplet import TripletLossWithMiner
-from oml.metrics.embeddings import EmbeddingMetricsDDP
+from oml.metrics.embeddings import EmbeddingMetrics
 from oml.miners.inbatch_all_tri import AllTripletsMiner
 from oml.models import ViTExtractor
 from oml.samplers.balance import BalanceSampler
@@ -25,16 +26,16 @@ df_train, df_val = download_mock_dataset(global_paths=True)
 extractor = ViTExtractor("vits16_dino", arch="vits16", normalise_features=False)
 
 # train
-optimizer = torch.optim.SGD(extractor.parameters(), lr=1e-6)
-train_dataset = DatasetWithLabels(df_train)
+optimizer = SGD(extractor.parameters(), lr=1e-6)
+train_dataset = ImageLabeledDataset(df_train)
 criterion = TripletLossWithMiner(margin=0.1, miner=AllTripletsMiner())
 batch_sampler = BalanceSampler(train_dataset.get_labels(), n_labels=2, n_instances=3)
-train_loader = torch.utils.data.DataLoader(train_dataset, batch_sampler=batch_sampler)
+train_loader = DataLoader(train_dataset, batch_sampler=batch_sampler)
 
 # val
-val_dataset = DatasetQueryGallery(df_val)
-val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=4)
-metric_callback = MetricValCallbackDDP(metric=EmbeddingMetricsDDP())  # DDP specific
+val_dataset = ImageQueryGalleryLabeledDataset(df_val)
+val_loader = DataLoader(val_dataset, batch_size=4)
+metric_callback = MetricValCallbackDDP(metric=EmbeddingMetrics(dataset=val_dataset))  # DDP specific
 
 # run
 pl_model = ExtractorModuleDDP(extractor=extractor, criterion=criterion, optimizer=optimizer,

docs/readme/examples_source/extractor/train_with_pml.md

@@ -6,7 +6,7 @@
 import torch
 from tqdm import tqdm
 
-from oml.datasets.base import DatasetWithLabels
+from oml.datasets import ImageLabeledDataset
 from oml.models import ViTExtractor
 from oml.samplers.balance import BalanceSampler
 from oml.utils.download_mock_dataset import download_mock_dataset
@@ -18,7 +18,7 @@ df_train, _ = download_mock_dataset(global_paths=True)
 extractor = ViTExtractor("vits16_dino", arch="vits16", normalise_features=False).train()
 optimizer = torch.optim.SGD(extractor.parameters(), lr=1e-6)
 
-train_dataset = DatasetWithLabels(df_train)
+train_dataset = ImageLabeledDataset(df_train)
 
 # PML specific
 # criterion = losses.TripletMarginLoss(margin=0.2, triplets_per_anchor="all")

docs/readme/examples_source/extractor/train_with_pml_advanced.md

@@ -6,7 +6,7 @@
 import torch
 from tqdm import tqdm
 
-from oml.datasets.base import DatasetWithLabels
+from oml.datasets import ImageLabeledDataset
 from oml.models import ViTExtractor
 from oml.samplers.balance import BalanceSampler
 from oml.utils.download_mock_dataset import download_mock_dataset
@@ -18,7 +18,7 @@ df_train, _ = download_mock_dataset(global_paths=True)
 extractor = ViTExtractor("vits16_dino", arch="vits16", normalise_features=False).train()
 optimizer = torch.optim.SGD(extractor.parameters(), lr=1e-6)
 
-train_dataset = DatasetWithLabels(df_train)
+train_dataset = ImageLabeledDataset(df_train)
 
 # PML specific
 distance = distances.LpDistance(p=2)