Simplify imports from losses, models, miners, and samplers in /oml module

Author: mimi030

oml/interfaces/datasets.py

@@ -13,7 +13,7 @@
     PAIR_1ST_KEY,
     PAIR_2ND_KEY,
 )
-from oml.samplers.balance import BalanceSampler  # noqa
+from oml.samplers import BalanceSampler  # noqa
 
 
 class IDatasetWithLabels(Dataset, ABC):

oml/lightning/modules/pairwise_postprocessing.py

@@ -9,7 +9,7 @@
 from oml.const import EMBEDDINGS_KEY, INPUT_TENSORS_KEY, LABELS_KEY
 from oml.interfaces.models import IFreezable, IPairwiseModel
 from oml.lightning.modules.ddp import ModuleDDP
-from oml.miners.pairs import PairsMiner
+from oml.miners import PairsMiner
 
 
 class PairwiseModule(pl.LightningModule):

oml/lightning/pipelines/train_postprocessor.py

@@ -28,7 +28,7 @@
     parse_scheduler_from_config,
 )
 from oml.metrics.embeddings import EmbeddingMetrics, EmbeddingMetricsDDP
-from oml.miners.pairs import PairsMiner
+from oml.miners import PairsMiner
 from oml.registry.models import get_extractor_by_cfg
 from oml.registry.optimizers import get_optimizer_by_cfg
 from oml.registry.postprocessors import get_postprocessor_by_cfg

oml/registry/losses.py

@@ -2,9 +2,9 @@
 
 from torch import nn
 
-from oml.losses.arcface import ArcFaceLoss, ArcFaceLossWithMLP
-from oml.losses.surrogate_precision import SurrogatePrecision
-from oml.losses.triplet import TripletLoss, TripletLossPlain, TripletLossWithMiner
+from oml.losses import ArcFaceLoss, ArcFaceLossWithMLP
+from oml.losses import SurrogatePrecision
+from oml.losses import TripletLoss, TripletLossPlain, TripletLossWithMiner
 from oml.registry.miners import get_miner_by_cfg
 from oml.utils.misc import TCfg, dictconfig_to_dict, remove_unused_kwargs
 

oml/registry/miners.py

@@ -1,12 +1,12 @@
 from typing import Any, Dict
 
 from oml.interfaces.miners import ITripletsMiner
-from oml.miners.cross_batch import TripletMinerWithMemory
-from oml.miners.inbatch_all_tri import AllTripletsMiner
-from oml.miners.inbatch_hard_cluster import HardClusterMiner
-from oml.miners.inbatch_hard_tri import HardTripletsMiner
-from oml.miners.inbatch_nhard_tri import NHardTripletsMiner
-from oml.miners.miner_with_bank import MinerWithBank
+from oml.miners import TripletMinerWithMemory
+from oml.miners import AllTripletsMiner
+from oml.miners import HardClusterMiner
+from oml.miners import HardTripletsMiner
+from oml.miners import NHardTripletsMiner
+from oml.miners import MinerWithBank
 from oml.utils.misc import TCfg, dictconfig_to_dict
 
 MINERS_REGISTRY = {

oml/registry/models.py

@@ -4,16 +4,16 @@
 from torch import nn
 
 from oml.interfaces.models import IExtractor, IPairwiseModel
-from oml.models.meta.projection import ExtractorWithMLP
-from oml.models.meta.siamese import (
+from oml.models import ExtractorWithMLP
+from oml.models import (
     ConcatSiamese,
     LinearTrivialDistanceSiamese,
     TrivialDistanceSiamese,
 )
-from oml.models.resnet.extractor import ResnetExtractor
-from oml.models.vit_clip.extractor import ViTCLIPExtractor
-from oml.models.vit_dino.extractor import ViTExtractor
-from oml.models.vit_unicom.extractor import ViTUnicomExtractor
+from oml.models import ResnetExtractor
+from oml.models import ViTCLIPExtractor
+from oml.models import ViTExtractor
+from oml.models import ViTUnicomExtractor
 from oml.utils.misc import TCfg, dictconfig_to_dict
 
 EXTRACTORS_REGISTRY = {

oml/registry/samplers.py

@@ -1,9 +1,9 @@
 from typing import Any, Dict
 
 from oml.interfaces.samplers import IBatchSampler
-from oml.samplers.balance import BalanceSampler
-from oml.samplers.category_balance import CategoryBalanceSampler
-from oml.samplers.distinct_category_balance import DistinctCategoryBalanceSampler
+from oml.samplers import BalanceSampler
+from oml.samplers import CategoryBalanceSampler
+from oml.samplers import DistinctCategoryBalanceSampler
 from oml.utils.misc import TCfg, dictconfig_to_dict, remove_unused_kwargs
 
 SAMPLERS_CATEGORIES_BASED = {

oml/registry/transforms.py

@@ -3,7 +3,7 @@
 
 import albumentations as albu
 
-import oml.models.vit_unicom.external.vision_transformer as unicom  # type: ignore
+import oml.models as unicom  # type: ignore
 from oml.transforms.images.albumentations import (
     get_augs_albu,
     get_normalisation_albu,