refactor: internal submodule

.coveragerc

@@ -3,9 +3,15 @@ omit =
   *_tmp.py
   zetta_utils/log.py
   zetta_utils/mazepa_layer_processing/*.py
+  zetta_utils/mazepa_layer_processing/**/*.py
   zetta_utils/tensor_mapping/*.py
+  zetta_utils/tensor_mapping/**/*.py
   zetta_utils/convnet/architecture/deprecated/*.py
+  zetta_utils/convnet/architecture/deprecated/**/*.py
   zetta_utils/viz/*.py
+  zetta_utils/viz/**/*.py
+  zetta_utils/internal/*.py
+  zetta_utils/internal/**/*.py
   zetta_utils/mazepa_addons/*.py
   zetta_utils/mazepa_addons/**/*.py
   zetta_utils/alignment/*.py

.github/workflows/testing.yaml

@@ -122,7 +122,7 @@ jobs:
         run: pylint ${{ steps.changed-py-files.outputs.all_changed_files }}
       - name: Run isort
         if: ${{ steps.changed-py-files.outputs.any_changed == 'true' }}
-        run: isort --check-only --df --verbose --profile black .
+        run: isort --check-only --df --verbose --om --profile black .
   mypy:
     strategy:
       matrix:

.gitmodules

@@ -1,3 +1,3 @@
-[submodule "zetta_utils/alignment"]
-	path = zetta_utils/alignment
-	url = [email protected]:ZettaAI/alignment.git
+[submodule "zetta_utils/internal"]
+	path = zetta_utils/internal
+	url = [email protected]:ZettaAI/internal.git

scripts/benchmark_rigidity_map.py

@@ -4,7 +4,7 @@
 import torch
 import torchfields
 
-from zetta_utils.alignment import field
+from zetta_utils.internal.alignment import field
 
 
 def rotation_tensor(degrees):

specs/nico/training/em_encoder/train/m3_m6_encoder_dict_c4.py

@@ -22,9 +22,7 @@
 
     EXP_VERSION = f"4.0.0_M3_M6_3px_C{CHANNELS}_lr{LR}_locality{LOCALITY_WEIGHT}_similarity{SIMILARITY_WEIGHT}_l1{L1_WEIGHT_START_VAL}-{L1_WEIGHT_END_VAL}_N1x4"
 
-    START_EXP_VERSION = (
-        f"4.4.0_M3_M6_C1_lr0.0002_locality1.0_similarity0.0_l10.05-0.12_N1x4"
-    )
+    START_EXP_VERSION = f"4.4.0_M3_M6_C1_lr0.0002_locality1.0_similarity0.0_l10.05-0.12_N1x4"
     MODEL_CKPT = None  # f"gs://zetta-research-nico/training_artifacts/general_coarsener_loss/{START_EXP_VERSION}/last.ckpt"
 
     BASE_PATH = "gs://zetta-research-nico/encoder/"

specs/nico/training/em_encoder/train/m3_m8_encoder_dict.py

@@ -22,9 +22,7 @@
 
     EXP_VERSION = f"1.0.0_M3_M8_C{CHANNELS}_lr{LR}_locality{LOCALITY_WEIGHT}_similarity{SIMILARITY_WEIGHT}_l1{L1_WEIGHT_START_VAL}-{L1_WEIGHT_END_VAL}_N1x4"
 
-    START_EXP_VERSION = (
-        f"1.2.0_M3_M7_C1_lr2e-05_locality1.0_similarity0.0_l10.12-0.12_N1x4"
-    )
+    START_EXP_VERSION = f"1.2.0_M3_M7_C1_lr2e-05_locality1.0_similarity0.0_l10.12-0.12_N1x4"
     MODEL_CKPT = None  # f"gs://zetta-research-nico/training_artifacts/general_coarsener_loss/{START_EXP_VERSION}/last.ckpt"
 
     BASE_PATH = "gs://zetta-research-nico/encoder/"

zetta_utils/__init__.py

@@ -35,7 +35,7 @@ def try_load_train_inference():  # pragma: no cover
 
 def try_load_submodules():  # pragma: no cover
     try:
-        from . import alignment
+        from . import internal
     except ImportError:
         ...
 
@@ -46,7 +46,6 @@ def load_inference_modules():
         convnet,
         mazepa,
         mazepa_layer_processing,
-        segmentation,
         tensor_ops,
         tensor_typing,
         tensor_mapping,

zetta_utils/api/v0.py

@@ -1,26 +1,6 @@
 # pylint: disable=unused-import
 
 
-from zetta_utils.alignment.aced_relaxation import (
-    compute_aced_loss_new,
-    get_aced_match_offsets,
-    get_aced_match_offsets_naive,
-    perform_aced_relaxation,
-)
-from zetta_utils.alignment.base_coarsener import BaseCoarsener
-from zetta_utils.alignment.base_encoder import BaseEncoder
-from zetta_utils.alignment.encoding_coarsener import EncodingCoarsener
-from zetta_utils.alignment.field import (
-    gen_biased_perlin_noise_field,
-    get_rigidity_map,
-    get_rigidity_map_zcxy,
-    invert_field,
-    invert_field_opti,
-    percentile,
-    profile_field2d_percentile,
-)
-from zetta_utils.alignment.misalignment_detector import MisalignmentDetector, naive_misd
-from zetta_utils.alignment.online_finetuner import align_with_online_finetuner
 from zetta_utils.augmentations.common import prob_aug
 from zetta_utils.augmentations.imgaug import imgaug_augment, imgaug_readproc
 from zetta_utils.augmentations.tensor import (
@@ -134,6 +114,37 @@
 from zetta_utils.geometry.bbox import BBox3D
 from zetta_utils.geometry.bbox_strider import BBoxStrider
 from zetta_utils.geometry.vec import Vec3D, allclose, is_int_vec, is_raw_vec3d, isclose
+from zetta_utils.internal.alignment.aced_relaxation import (
+    compute_aced_loss_new,
+    get_aced_match_offsets,
+    get_aced_match_offsets_naive,
+    perform_aced_relaxation,
+)
+from zetta_utils.internal.alignment.base_coarsener import BaseCoarsener
+from zetta_utils.internal.alignment.base_encoder import BaseEncoder
+from zetta_utils.internal.alignment.encoding_coarsener import EncodingCoarsener
+from zetta_utils.internal.alignment.field import (
+    gen_biased_perlin_noise_field,
+    get_rigidity_map,
+    get_rigidity_map_zcxy,
+    invert_field,
+    invert_field_opti,
+    percentile,
+    profile_field2d_percentile,
+)
+from zetta_utils.internal.alignment.misalignment_detector import (
+    MisalignmentDetector,
+    naive_misd,
+)
+from zetta_utils.internal.alignment.online_finetuner import align_with_online_finetuner
+from zetta_utils.internal.regimes.alignment.base_encoder import BaseEncoderRegime
+from zetta_utils.internal.regimes.alignment.misalignment_detector_aced import (
+    MisalignmentDetectorAcedRegime,
+)
+from zetta_utils.internal.regimes.common import is_2d_image, log_results
+from zetta_utils.internal.regimes.naive_supervised import NaiveSupervisedRegime
+from zetta_utils.internal.regimes.noop import NoOpRegime
+from zetta_utils.internal.segmentation.inference import run_affinities_inference_onnx
 from zetta_utils.layer.backend_base import Backend
 from zetta_utils.layer.db_layer.backend import DBBackend
 from zetta_utils.layer.db_layer.build import build_db_layer
@@ -352,7 +363,6 @@
     read_remote_annotations,
     write_remote_annotations,
 )
-from zetta_utils.segmentation.inference import run_affinities_inference_onnx
 from zetta_utils.tensor_mapping.tensor_mapping import TensorMapping
 from zetta_utils.tensor_ops.common import (
     add,
@@ -382,7 +392,8 @@
     kornia_opening,
     skip_on_empty_data,
 )
-from zetta_utils.tensor_ops.transform import get_affine_field
+
+# from zetta_utils.tensor_ops.generators import get_affine_field
 from zetta_utils.training.datasets.joint_dataset import JointDataset
 from zetta_utils.training.datasets.layer_dataset import LayerDataset
 from zetta_utils.training.datasets.sample_indexers.base import SampleIndexer
@@ -395,17 +406,6 @@
 from zetta_utils.training.datasets.sample_indexers.volumetric_strided_indexer import (
     VolumetricStridedIndexer,
 )
-from zetta_utils.training.lightning.regimes.alignment.base_encoder import (
-    BaseEncoderRegime,
-)
-from zetta_utils.training.lightning.regimes.alignment.misalignment_detector_aced import (
-    MisalignmentDetectorAcedRegime,
-)
-from zetta_utils.training.lightning.regimes.common import is_2d_image, log_results
-from zetta_utils.training.lightning.regimes.naive_supervised import (
-    NaiveSupervisedRegime,
-)
-from zetta_utils.training.lightning.regimes.noop import NoOpRegime
 from zetta_utils.training.lightning.train import lightning_train
 from zetta_utils.training.lightning.trainers.default import (
     ConfigureTraceCallback,

zetta_utils/mazepa_layer_processing/__init__.py

@@ -6,8 +6,8 @@
     ChunkableOpProtocol,
     VolumetricOpProtocol,
 )
-from . import alignment
-from . import segmentation
+from . import segmentation, alignment
+
 from .common import (
     ChunkedApplyFlowSchema,
     CallableOperation,

zetta_utils/mazepa_layer_processing/alignment/aced_relaxation_flow.py

@@ -5,8 +5,9 @@
 import attrs
 import torch
 
-from zetta_utils import alignment, builder, mazepa, tensor_ops
+from zetta_utils import builder, mazepa, tensor_ops
 from zetta_utils.geometry import BBox3D, Vec3D
+from zetta_utils.internal import alignment
 from zetta_utils.layer.volumetric import (
     VolumetricIndex,
     VolumetricIndexChunker,

zetta_utils/mazepa_layer_processing/alignment/common.py

@@ -4,8 +4,8 @@
 
 import torch
 
-from zetta_utils import alignment
 from zetta_utils.geometry import Vec3D
+from zetta_utils.internal import alignment
 from zetta_utils.layer.volumetric import VolumetricIndex, VolumetricLayer
 
 

zetta_utils/training/lightning/__init__.py

@@ -1 +1 @@
-from . import regimes, train, trainers
+from . import train, trainers