From b21af13a70ad63904a988da051363b4b7b7fb009 Mon Sep 17 00:00:00 2001 From: Sergiy Popovych Date: Fri, 2 Feb 2024 03:01:51 +0000 Subject: [PATCH] refactor: internal submodule --- .coveragerc | 6 ++ .github/workflows/testing.yaml | 2 +- .gitmodules | 6 +- scripts/benchmark_rigidity_map.py | 2 +- .../em_encoder/train/m3_m6_encoder_dict_c4.py | 4 +- .../em_encoder/train/m3_m8_encoder_dict.py | 4 +- tmp.py | 0 zetta_utils/__init__.py | 3 +- zetta_utils/api/v0.py | 66 +++++++++---------- zetta_utils/internal | 1 + .../mazepa_layer_processing/__init__.py | 4 +- .../alignment/aced_relaxation_flow.py | 3 +- .../alignment/common.py | 2 +- zetta_utils/training/lightning/__init__.py | 2 +- 14 files changed, 54 insertions(+), 51 deletions(-) create mode 100644 tmp.py create mode 160000 zetta_utils/internal diff --git a/.coveragerc b/.coveragerc index 0e7de3e75..c6ba6ddc7 100644 --- a/.coveragerc +++ b/.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 diff --git a/.github/workflows/testing.yaml b/.github/workflows/testing.yaml index a021139a5..555aa0c2c 100644 --- a/.github/workflows/testing.yaml +++ b/.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: diff --git a/.gitmodules b/.gitmodules index 151433cc5..44216a04b 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,3 @@ -[submodule "zetta_utils/alignment"] - path = zetta_utils/alignment - url = git@github.com:ZettaAI/alignment.git +[submodule "zetta_utils/internal"] + path = zetta_utils/internal + url = git@github.com:ZettaAI/internal.git diff --git a/scripts/benchmark_rigidity_map.py b/scripts/benchmark_rigidity_map.py index 10dc9bcc4..b7f1dc5e7 100644 --- a/scripts/benchmark_rigidity_map.py +++ b/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): diff --git a/specs/nico/training/em_encoder/train/m3_m6_encoder_dict_c4.py b/specs/nico/training/em_encoder/train/m3_m6_encoder_dict_c4.py index 95b50fe82..6a414fbf9 100644 --- a/specs/nico/training/em_encoder/train/m3_m6_encoder_dict_c4.py +++ b/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/" diff --git a/specs/nico/training/em_encoder/train/m3_m8_encoder_dict.py b/specs/nico/training/em_encoder/train/m3_m8_encoder_dict.py index 5a11d757f..92feafcc9 100644 --- a/specs/nico/training/em_encoder/train/m3_m8_encoder_dict.py +++ b/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/" diff --git a/tmp.py b/tmp.py new file mode 100644 index 000000000..e69de29bb diff --git a/zetta_utils/__init__.py b/zetta_utils/__init__.py index 3dca33a58..1513ae044 100644 --- a/zetta_utils/__init__.py +++ b/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, diff --git a/zetta_utils/api/v0.py b/zetta_utils/api/v0.py index 4f9fa907f..47b26d2d1 100644 --- a/zetta_utils/api/v0.py +++ b/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, diff --git a/zetta_utils/internal b/zetta_utils/internal new file mode 160000 index 000000000..beffa7446 --- /dev/null +++ b/zetta_utils/internal @@ -0,0 +1 @@ +Subproject commit beffa7446d8cdb0bfa67e9c7080573dfe966dcd7 diff --git a/zetta_utils/mazepa_layer_processing/__init__.py b/zetta_utils/mazepa_layer_processing/__init__.py index fa28f23f0..88e5a99b1 100644 --- a/zetta_utils/mazepa_layer_processing/__init__.py +++ b/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, diff --git a/zetta_utils/mazepa_layer_processing/alignment/aced_relaxation_flow.py b/zetta_utils/mazepa_layer_processing/alignment/aced_relaxation_flow.py index fb207a316..26181eb44 100644 --- a/zetta_utils/mazepa_layer_processing/alignment/aced_relaxation_flow.py +++ b/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, diff --git a/zetta_utils/mazepa_layer_processing/alignment/common.py b/zetta_utils/mazepa_layer_processing/alignment/common.py index 82ee69a40..342c00fe5 100644 --- a/zetta_utils/mazepa_layer_processing/alignment/common.py +++ b/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 diff --git a/zetta_utils/training/lightning/__init__.py b/zetta_utils/training/lightning/__init__.py index 7142ef097..76e8338b9 100644 --- a/zetta_utils/training/lightning/__init__.py +++ b/zetta_utils/training/lightning/__init__.py @@ -1 +1 @@ -from . import regimes, train, trainers +from . import train, trainers