some baseline code for the repo

.vscode/settings.json

@@ -0,0 +1,5 @@
+{
+    "flake8.args": [
+        "--max-line-length=120"
+    ]
+}

requirements.txt

@@ -0,0 +1,212 @@
+accelerate==0.24.1
+aiohttp==3.8.6
+aiohttp-retry==2.8.3
+aiosignal==1.3.1
+albumentations==1.3.0
+amqp==5.1.1
+analytics-python==1.4.post1
+antlr4-python3-runtime==4.9.3
+appdirs==1.4.4
+argcomplete==3.1.2
+asciitree==0.3.3
+async-generator==1.10
+async-timeout==4.0.3
+asyncssh==2.14.0
+atpublic==4.0
+attrs==23.1.0
+azure-core==1.29.4
+azure-storage-blob==12.18.3
+backoff==1.10.0
+bcrypt==4.0.1
+billiard==4.1.0
+boto3==1.28.64
+botocore==1.31.64
+build==1.0.3
+cachetools==5.3.1
+celery==5.3.4
+certifi==2023.7.22
+cffi==1.16.0
+charset-normalizer==3.3.0
+click==8.1.7
+click-didyoumean==0.3.0
+click-plugins==1.1.1
+click-repl==0.3.0
+cloudpickle==2.2.1
+colorama==0.4.6
+configobj==5.0.8
+contourpy==1.1.1
+crc32c==2.3.post0
+cryptography==41.0.4
+cycler==0.12.1
+Cython==0.29.36
+databricks-cli==0.18.0
+db-dtypes==1.1.1
+decorator==5.1.1
+determined==0.23.3
+dictdiffer==0.9.0
+diskcache==5.6.3
+distro==1.8.0
+docker==6.1.3
+dpath==2.1.6
+dulwich==0.21.6
+dvc==3.26.2
+dvc-data==2.18.1
+dvc-gs==2.22.1
+dvc-http==2.30.2
+dvc-objects==1.0.1
+dvc-render==0.6.0
+dvc-studio-client==0.15.0
+dvc-task==0.3.0
+entrypoints==0.4
+fasteners==0.19
+filelock==3.12.4
+flatten-dict==0.4.2
+flufl.lock==7.1.1
+fonttools==4.43.1
+frozenlist==1.4.0
+fsspec==2023.9.2
+funcy==2.0
+gcsfs==2023.9.2
+gitdb==4.0.10
+GitPython==3.1.38
+google-api-core==2.12.0
+google-api-python-client==2.103.0
+google-auth==2.23.3
+google-auth-httplib2==0.1.1
+google-auth-oauthlib==1.0.0
+google-cloud==0.34.0
+google-cloud-bigquery==3.12.0
+google-cloud-core==2.3.3
+google-cloud-storage==2.12.0
+google-crc32c==1.5.0
+google-resumable-media==2.6.0
+googleapis-common-protos==1.61.0
+grandalf==0.8
+grpcio==1.59.0
+grpcio-status==1.48.2
+gto==1.4.0
+httplib2==0.22.0
+huggingface-hub==0.18.0
+hydra-core==1.3.2
+idna==3.4
+importlib-metadata==6.8.0
+isodate==0.6.1
+iterative-telemetry==0.0.8
+Jinja2==3.1.2
+jmespath==1.0.1
+joblib==1.3.2
+kiwisolver==1.4.5
+kombu==5.3.2
+lightning-fabric==2.1.0
+lightning-utilities==0.9.0
+lmdb==1.4.1
+lomond==0.3.3
+mahotas==1.4.13
+markdown-it-py==3.0.0
+MarkupSafe==2.1.3
+matplotlib==3.8.0
+mdurl==0.1.2
+mlflow-skinny==2.6.0
+monotonic==1.6
+mpmath==0.19
+multidict==6.0.4
+networkx==3.1
+numcodecs==0.12.0
+numpy==1.26.1
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.18.1
+nvidia-nvjitlink-cu12==12.2.140
+nvidia-nvtx-cu12==12.1.105
+oauthlib==3.2.2
+omegaconf==2.3.0
+orjson==3.9.9
+packaging==21.3
+pandas==2.1.1
+paramiko==3.3.1
+pathspec==0.11.2
+peft==0.5.0
+Pillow==10.1.0
+pip-tools==7.3.0
+platformdirs==3.11.0
+prompt-toolkit==3.0.39
+proto-plus==1.22.3
+protobuf==3.20.3
+psutil==5.9.6
+pyarrow==13.0.0
+pyasn1==0.5.0
+pyasn1-modules==0.3.0
+pycparser==2.21
+pydantic==1.10.13
+pydot==1.4.2
+pygit2==1.13.1
+Pygments==2.16.1
+pygtrie==2.5.0
+PyJWT==2.8.0
+PyNaCl==1.5.0
+pyOpenSSL==23.2.0
+pyparsing==3.0.9
+pyproject_hooks==1.0.0
+python-dateutil==2.8.2
+pytorch-lightning==2.1.0
+pytz==2023.3.post1
+PyYAML==6.0.1
+pyzmq==25.1.1
+regex==2023.10.3
+requests==2.31.0
+requests-oauthlib==1.3.1
+rich==13.6.0
+rsa==4.9
+ruamel.yaml==0.17.35
+ruamel.yaml.clib==0.2.8
+s3transfer==0.7.0
+safetensors==0.4.0
+scikit-learn==1.3.1
+scipy==1.11.3
+scmrepo==1.4.0
+semver==3.0.2
+shortuuid==1.0.11
+shtab==1.6.4
+six==1.16.0
+smmap==5.0.1
+sqlparse==0.4.4
+sqltrie==0.8.0
+sympy==1.12
+tabulate==0.9.0
+tensorboardX==2.6.2.2
+termcolor==2.3.0
+threadpoolctl==3.2.0
+timm==0.9.7
+tokenizers==0.15.0
+tomli==2.0.1
+tomlkit==0.12.1
+toolz==0.12.0
+torch==2.1.0+cu121
+torchmetrics==1.2.0
+torchvision==0.16.0+cu121
+tqdm==4.66.1
+transformers==4.35.2
+triton==2.1.0
+typer==0.9.0
+typing_extensions==4.8.0
+tzdata==2023.3
+uritemplate==4.1.1
+urllib3==1.26.17
+vine==5.0.0
+voluptuous==0.13.1
+wcwidth==0.2.8
+websocket-client==1.6.4
+websockets==11.0.3
+xformers==0.0.22.post7
+yarl==1.9.2
+yogadl==0.1.4
+zarr==2.16.1
+zc.lockfile==3.0.post1
+zipp==3.17.0

vit_encoder.py

@@ -0,0 +1,60 @@
+from typing import Dict
+
+import timm.models.vision_transformer as vit
+import torch
+
+
+def build_imagenet_baselines() -> Dict[str, torch.jit.ScriptModule]:
+    """This returns the prepped imagenet encoders from timm, not bad for microscopy data."""
+    vit_backbones = [
+        _make_vit(vit.vit_small_patch16_384),
+        _make_vit(vit.vit_base_patch16_384),
+        _make_vit(vit.vit_base_patch8_224),
+        _make_vit(vit.vit_large_patch16_384),
+    ]
+    model_names = [
+        "vit_small_patch16_384",
+        "vit_base_patch16_384",
+        "vit_base_patch8_224",
+        "vit_large_patch16_384",
+    ]
+    imagenet_encoders = list(map(_make_torchscripted_encoder, vit_backbones))
+    return {name: model for name, model in zip(model_names, imagenet_encoders)}
+
+
+def _make_torchscripted_encoder(vit_backbone) -> torch.jit.ScriptModule:
+    dummy_input = torch.testing.make_tensor(
+        (2, 6, 256, 256),
+        low=0,
+        high=255,
+        dtype=torch.uint8,
+        device=torch.device("cpu"),
+    )
+    encoder = torch.nn.Sequential(
+        Normalizer(),
+        torch.nn.LazyInstanceNorm2d(
+            affine=False, track_running_stats=False
+        ),  # this module performs self-standardization, very important
+        vit_backbone,
+    ).to(device="cpu")
+    _ = encoder(dummy_input)  # get those lazy modules built
+    return torch.jit.freeze(torch.jit.script(encoder.eval()))
+
+
+def _make_vit(constructor):
+    return constructor(
+        pretrained=True,  # download imagenet weights
+        img_size=256,  # 256x256 crops
+        in_chans=6,  # we expect 6-channel microscopy images
+        num_classes=0,
+        fc_norm=None,
+        class_token=True,
+        global_pool="avg",  # minimal perf diff btwn "cls" and "avg"
+    )
+
+
+class Normalizer(torch.nn.Module):
+    def forward(self, pixels: torch.Tensor) -> torch.Tensor:
+        pixels = pixels.float()
+        pixels /= 255.0
+        return pixels