From 8028ecf6a983914c9697959979b49c718fc3f70f Mon Sep 17 00:00:00 2001 From: Jean-Guillaume Durand Date: Mon, 12 Aug 2024 12:43:58 -0700 Subject: [PATCH 1/2] Added packaging files --- satclip/.gitignore => .gitignore | 4 +++- pyproject.toml | 39 ++++++++++++++++++++++++++++++++ satclip/load.py | 4 ++-- satclip/load_lightweight.py | 10 ++++---- satclip/location_encoder.py | 10 ++++---- satclip/main.py | 7 +++--- satclip/model.py | 30 ++++++++++++------------ 7 files changed, 72 insertions(+), 32 deletions(-) rename satclip/.gitignore => .gitignore (55%) create mode 100644 pyproject.toml diff --git a/satclip/.gitignore b/.gitignore similarity index 55% rename from satclip/.gitignore rename to .gitignore index 52804a2..e3a142f 100644 --- a/satclip/.gitignore +++ b/.gitignore @@ -2,4 +2,6 @@ satclip_logs/ # Python -__pycache__/ \ No newline at end of file +__pycache__/ +*.egg-info/ +build/ diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..5f7bbef --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,39 @@ +[build-system] +requires = ["setuptools >= 61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "satclip" +description = "A global, general-purpose geographic location encoder" +version = "0.0.1" +authors = [ + {name="Konstantin Klemmer"}, + {name="Esther Rolf"}, + {name="Caleb Robinson"}, + {name="Lester Mackey"}, + {name="Marc Rußwurm"}, +] +readme = "README.md" +license = {file = "LICENSE"} +requires-python = ">=3.9" +classifiers = [ + "Programming Language :: Python :: 3", + "Operating System :: OS Independent", + "License :: OSI Approved :: MIT License", +] + +dependencies = [ + "albumentations", + "lightning == 2.2.2", + "pandas", + "rasterio >= 1.3.10", + "torchgeo >= 0.5", # Forces Python 3.9+ +] + +[project.urls] +Homepage = "https://github.com/microsoft/satclip" +Repository = "https://github.com/microsoft/satclip.git" +Issues = "https://github.com/microsoft/satclip/issues" + +[tool.setuptools.packages.find] +include = ["satclip", "satclip.*"] diff --git a/satclip/load.py b/satclip/load.py index ac975de..5c8448f 100644 --- a/satclip/load.py +++ b/satclip/load.py @@ -1,4 +1,4 @@ -from main import * +from .main import * def get_satclip(ckpt_path, device, return_all=False): ckpt = torch.load(ckpt_path,map_location=device) @@ -15,4 +15,4 @@ def get_satclip(ckpt_path, device, return_all=False): if return_all: return geo_model else: - return geo_model.location \ No newline at end of file + return geo_model.location diff --git a/satclip/load_lightweight.py b/satclip/load_lightweight.py index 6dfc828..2ec47d3 100644 --- a/satclip/load_lightweight.py +++ b/satclip/load_lightweight.py @@ -1,5 +1,6 @@ import torch -from location_encoder import get_neural_network, get_positional_encoding, LocationEncoder + +from .location_encoder import get_neural_network, get_positional_encoding, LocationEncoder def get_satclip_loc_encoder(ckpt_path, device): @@ -14,7 +15,7 @@ def get_satclip_loc_encoder(ckpt_path, device): hp['max_radius'], hp['frequency_num'] ) - + nnet = get_neural_network( hp['pe_type'], posenc.embedding_dim, @@ -25,12 +26,11 @@ def get_satclip_loc_encoder(ckpt_path, device): # only load nnet params from state dict state_dict = ckpt['state_dict'] - state_dict = {k[k.index('nnet'):]:state_dict[k] + state_dict = {k[k.index('nnet'):]:state_dict[k] for k in state_dict.keys() if 'nnet' in k} - + loc_encoder = LocationEncoder(posenc, nnet).double() loc_encoder.load_state_dict(state_dict) loc_encoder.eval() return loc_encoder - \ No newline at end of file diff --git a/satclip/location_encoder.py b/satclip/location_encoder.py index 7bf1617..1911ac9 100644 --- a/satclip/location_encoder.py +++ b/satclip/location_encoder.py @@ -1,11 +1,11 @@ -from torch import nn, optim import math + import torch import torch.nn.functional as F from einops import rearrange -import numpy as np -from datetime import datetime -import positional_encoding as PE +from torch import nn + +from . import positional_encoding as PE """ FCNet @@ -110,7 +110,7 @@ def forward(self, x, mods = None): x *= rearrange(mod, 'd -> () d') return self.last_layer(x) - + class Sine(nn.Module): def __init__(self, w0 = 1.): super().__init__() diff --git a/satclip/main.py b/satclip/main.py index 03d1daf..a5134ab 100644 --- a/satclip/main.py +++ b/satclip/main.py @@ -3,10 +3,11 @@ import lightning.pytorch import torch -from datamodules.s2geo_dataset import S2GeoDataModule from lightning.pytorch.cli import LightningCLI -from loss import SatCLIPLoss -from model import SatCLIP + +from .datamodules.s2geo_dataset import S2GeoDataModule +from .loss import SatCLIPLoss +from .model import SatCLIP torch.set_float32_matmul_precision('high') diff --git a/satclip/model.py b/satclip/model.py index b75a129..26b400f 100644 --- a/satclip/model.py +++ b/satclip/model.py @@ -1,17 +1,15 @@ from collections import OrderedDict -from typing import Tuple, Union, Optional +from typing import Tuple, Union import numpy as np import torch import torch.nn.functional as F from torch import nn -import math import timm -import torchgeo.models from torchgeo.models import ResNet18_Weights, ResNet50_Weights, ViTSmall16_Weights -from location_encoder import get_positional_encoding, get_neural_network, LocationEncoder -from datamodules.s2geo_dataset import S2Geo + +from .location_encoder import get_positional_encoding, get_neural_network, LocationEncoder class Bottleneck(nn.Module): expansion = 4 @@ -257,12 +255,12 @@ def __init__(self, # location le_type: str, pe_type: str, - frequency_num: int, - max_radius: int, + frequency_num: int, + max_radius: int, min_radius: int, harmonics_calculation: str, - legendre_polys: int=10, - sh_embedding_dims: int=16, + legendre_polys: int=10, + sh_embedding_dims: int=16, ffn: bool=True, num_hidden_layers: int=2, capacity: int=256, @@ -270,7 +268,7 @@ def __init__(self, **kwargs ): super().__init__() - + if isinstance(vision_layers, (tuple, list)): print('using modified resnet') vision_heads = vision_width * 32 // 64 @@ -282,7 +280,7 @@ def __init__(self, width=vision_width, in_channels=in_channels ) - + elif vision_layers == 'moco_resnet18': print('using pretrained moco resnet18') weights = ResNet18_Weights.SENTINEL2_ALL_MOCO @@ -300,7 +298,7 @@ def __init__(self, self.visual.load_state_dict(weights.get_state_dict(progress=True), strict=False) self.visual.requires_grad_(False) self.visual.fc.requires_grad_(True) - + elif vision_layers == 'moco_vit16': print('using pretrained moco vit16') weights = ViTSmall16_Weights.SENTINEL2_ALL_MOCO @@ -322,13 +320,13 @@ def __init__(self, output_dim=embed_dim, in_channels=in_channels ) - + self.posenc = get_positional_encoding(name=le_type, harmonics_calculation=harmonics_calculation, legendre_polys=legendre_polys, min_radius=min_radius, max_radius=max_radius, frequency_num=frequency_num).double() self.nnet = get_neural_network(name=pe_type, input_dim=self.posenc.embedding_dim, num_classes=embed_dim, dim_hidden=capacity, num_layers=num_hidden_layers).double() - self.location = LocationEncoder(self.posenc, + self.location = LocationEncoder(self.posenc, self.nnet ).double() - + self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07)) self.initialize_parameters() @@ -362,7 +360,7 @@ def encode_location(self, coords): def forward(self, image, coords): - image_features = self.encode_image(image) + image_features = self.encode_image(image) location_features = self.encode_location(coords).float() # normalized features image_features = image_features / image_features.norm(dim=1, keepdim=True) From 837fc71caf3d8c0b1b6faff34bfc8366839ba44a Mon Sep 17 00:00:00 2001 From: Konstantin Klemmer Date: Fri, 12 Jun 2026 19:25:46 +0000 Subject: [PATCH 2/2] Polish packaging: docs, deps, configs, entry point, CI Build on the packaging work with the follow-ups needed to make the package usable end-to-end: - docs/notebooks: migrate README quickstart and notebooks (A01/A02/B01/B02) off 'git clone + sys.path' to 'pip install' + 'from satclip... import ...'; add an Installation section and run training via 'python -m satclip.main' - pyproject: list direct deps explicitly (torch, numpy, sympy, timm, einops, matplotlib, torchvision, huggingface_hub), relax 'lightning' pin to >=2.2,<3, SPDX license, add 'examples'/'dev' extras and a 'satclip-train' entry point, ship configs/*.yaml as package data - main.py: resolve the default config relative to the package and guard the CUDA device check so 'python -m satclip.main' works from any directory / on CPU - __init__.py: expose a clean public API (SatCLIP, LocationEncoder, SatCLIPLoss, get_satclip) instead of wildcard re-exports - add a CI workflow and import/forward-pass smoke tests - ignore dist/ and .pytest_cache/ Verified: wheel builds, bundles configs/default.yaml, registers the satclip-train entry point, and emits License-Expression: MIT. --- .github/workflows/ci.yml | 39 ++++++++++ .gitignore | 2 + README.md | 26 +++++-- notebooks/A01_Simple_SatCLIP_Usage.ipynb | 8 +-- .../A02_SatCLIP_Hugging_Face_Usage.ipynb | 10 +-- ...1_Example_Air_Temperature_Prediction.ipynb | 8 +-- .../B02_Example_Image_Localization.ipynb | 8 +-- pyproject.toml | 43 +++++++++-- satclip/__init__.py | 16 +++-- satclip/main.py | 16 +++-- tests/test_smoke.py | 71 +++++++++++++++++++ 11 files changed, 199 insertions(+), 48 deletions(-) create mode 100644 .github/workflows/ci.yml create mode 100644 tests/test_smoke.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..7b1caaa --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,39 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + +jobs: + test: + name: Build & test (Python ${{ matrix.python-version }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.9", "3.11"] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install torch (CPU) + # Install the CPU-only torch build first to keep CI lightweight and avoid + # pulling large CUDA wheels via transitive dependencies. + run: | + python -m pip install --upgrade pip + pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu + + - name: Install package + run: pip install -e ".[dev,examples]" + + - name: Import smoke test + run: python -c "import satclip; print('satclip public API:', satclip.__all__)" + + - name: Run tests + run: pytest -q diff --git a/.gitignore b/.gitignore index e3a142f..e863259 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,5 @@ satclip_logs/ __pycache__/ *.egg-info/ build/ +dist/ +.pytest_cache/ diff --git a/README.md b/README.md index 31ac775..ae6ed02 100644 --- a/README.md +++ b/README.md @@ -8,13 +8,27 @@ SatCLIP trains location and image encoders via contrastive learning, by matching images to their corresponding locations. This is analogous to the CLIP approach, which matches images to their corresponding text. Through this process, the location encoder learns characteristics of a location, as represented by satellite imagery. For more details, check out our [paper](https://arxiv.org/abs/2311.17179). +## Installation + +SatCLIP is packaged as a Python module. Install it directly from GitHub: +```bash +pip install git+https://github.com/microsoft/satclip.git +``` +Or, for a local/editable install (e.g. for training or development): +```bash +git clone https://github.com/microsoft/satclip.git +cd satclip +pip install -e . +``` +Once installed, `satclip` can be imported from any directory — no need to set `PYTHONPATH`. + ## Overview Usage of SatCLIP is simple: ```python -from model import * -from location_encoder import * +import torch +from satclip.model import SatCLIP model = SatCLIP( embed_dim=512, @@ -48,11 +62,11 @@ mkdir -p images for f in data/shard-*.tar; do tar -xf "$f" -C images; done ``` -Now, to train **SatCLIP** models, set the paths correctly (point `data.data_dir` in `satclip/configs/default.yaml` to this dataset directory), adapt training configs in `satclip/configs/default.yaml` and train SatCLIP by running: +Now, to train **SatCLIP** models, set the paths correctly (point `data.data_dir` in `satclip/configs/default.yaml` to this dataset directory), adapt training configs in `satclip/configs/default.yaml` and train SatCLIP by running the training module from the repository root: ```bash -cd satclip -python main.py +python -m satclip.main ``` +This requires an editable install (`pip install -e .`, see [Installation](#installation)). You can point to a custom config with `python -m satclip.main --config path/to/config.yaml`. ### Use of the S2-100K dataset @@ -96,7 +110,7 @@ We provide six pretrained SatCLIP models, trained with different vision encoders Usage of pretrained models is simple. Simply specify the SatCLIP model you want to access, e.g. `satclip-vit16-l40`: ```python from huggingface_hub import hf_hub_download -from load import get_satclip +from satclip.load import get_satclip import torch device = "cuda" if torch.cuda.is_available() else "cpu" diff --git a/notebooks/A01_Simple_SatCLIP_Usage.ipynb b/notebooks/A01_Simple_SatCLIP_Usage.ipynb index f3f062c..5b6bd3a 100644 --- a/notebooks/A01_Simple_SatCLIP_Usage.ipynb +++ b/notebooks/A01_Simple_SatCLIP_Usage.ipynb @@ -51,8 +51,7 @@ } ], "source": [ - "!rm -r sample_data .config # Empty current directory\n", - "!git clone https://github.com/microsoft/satclip.git . # Clone SatCLIP repository" + "!pip install git+https://github.com/microsoft/satclip.git" ] }, { @@ -169,11 +168,8 @@ { "cell_type": "code", "source": [ - "import sys\n", - "sys.path.append('./satclip')\n", - "\n", "import torch\n", - "from load import get_satclip" + "from satclip.load import get_satclip" ], "metadata": { "id": "grEIwoFjoHvu" diff --git a/notebooks/A02_SatCLIP_Hugging_Face_Usage.ipynb b/notebooks/A02_SatCLIP_Hugging_Face_Usage.ipynb index cf1c207..6c39233 100644 --- a/notebooks/A02_SatCLIP_Hugging_Face_Usage.ipynb +++ b/notebooks/A02_SatCLIP_Hugging_Face_Usage.ipynb @@ -397,8 +397,7 @@ } ], "source": [ - "!rm -r sample_data .config # Empty current directory\n", - "!git clone https://github.com/microsoft/satclip.git . # Clone SatCLIP repository" + "!pip install git+https://github.com/microsoft/satclip.git" ] }, { @@ -473,11 +472,8 @@ { "cell_type": "code", "source": [ - "import sys\n", - "sys.path.append('./satclip')\n", - "\n", "import torch\n", - "from load import get_satclip" + "from satclip.load import get_satclip" ], "metadata": { "id": "grEIwoFjoHvu" @@ -498,7 +494,7 @@ "cell_type": "code", "source": [ "from huggingface_hub import hf_hub_download\n", - "from load import get_satclip\n", + "from satclip.load import get_satclip\n", "import torch\n", "\n", "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", diff --git a/notebooks/B01_Example_Air_Temperature_Prediction.ipynb b/notebooks/B01_Example_Air_Temperature_Prediction.ipynb index ba892c3..1ece3a4 100644 --- a/notebooks/B01_Example_Air_Temperature_Prediction.ipynb +++ b/notebooks/B01_Example_Air_Temperature_Prediction.ipynb @@ -55,8 +55,7 @@ } ], "source": [ - "!rm -r sample_data .config # Empty current directory\n", - "!git clone https://github.com/microsoft/satclip.git . # Clone SatCLIP repository" + "!pip install git+https://github.com/microsoft/satclip.git" ] }, { @@ -168,11 +167,8 @@ { "cell_type": "code", "source": [ - "import sys\n", - "sys.path.append('./satclip')\n", - "\n", "import torch\n", - "from load import get_satclip\n", + "from satclip.load import get_satclip\n", "\n", "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Automatically select device" ], diff --git a/notebooks/B02_Example_Image_Localization.ipynb b/notebooks/B02_Example_Image_Localization.ipynb index 8a35b76..32e74a0 100644 --- a/notebooks/B02_Example_Image_Localization.ipynb +++ b/notebooks/B02_Example_Image_Localization.ipynb @@ -399,8 +399,7 @@ } ], "source": [ - "!rm -r sample_data .config # Empty current directory\n", - "!git clone https://github.com/microsoft/satclip.git . # Clone SatCLIP repository" + "!pip install git+https://github.com/microsoft/satclip.git" ] }, { @@ -485,10 +484,7 @@ { "cell_type": "code", "source": [ - "import sys\n", - "sys.path.append('./satclip')\n", - "\n", - "from load import get_satclip\n", + "from satclip.load import get_satclip\n", "from huggingface_hub import hf_hub_download\n", "import torch\n", "from sklearn.metrics.pairwise import cosine_similarity\n", diff --git a/pyproject.toml b/pyproject.toml index 5f7bbef..30d98b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["setuptools >= 61.0"] +requires = ["setuptools >= 77.0"] build-backend = "setuptools.build_meta" [project] @@ -14,22 +14,49 @@ authors = [ {name="Marc Rußwurm"}, ] readme = "README.md" -license = {file = "LICENSE"} +license = "MIT" +license-files = ["LICENSE"] requires-python = ">=3.9" classifiers = [ "Programming Language :: Python :: 3", "Operating System :: OS Independent", - "License :: OSI Approved :: MIT License", ] +# Direct (first-order) imports of the `satclip` package. Versions of the heavier +# scientific stack (numpy, sympy, torch, ...) are largely constrained by +# torchgeo/lightning, but are listed explicitly so the package does not silently +# rely on transitive dependencies. dependencies = [ "albumentations", - "lightning == 2.2.2", + "einops", + "huggingface_hub", + "lightning >=2.2,<3", + "matplotlib", + "numpy", "pandas", - "rasterio >= 1.3.10", - "torchgeo >= 0.5", # Forces Python 3.9+ + "rasterio >=1.3.10", + "sympy", + "timm", + "torch >=1.13", + "torchgeo >=0.5", # Forces Python 3.9+ + "torchvision", ] +[project.optional-dependencies] +# Extra dependencies used by the example notebooks only. +examples = [ + "scikit-learn", +] +# Development / CI tooling. +dev = [ + "build", + "pytest", + "ruff", +] + +[project.scripts] +satclip-train = "satclip.main:cli_main" + [project.urls] Homepage = "https://github.com/microsoft/satclip" Repository = "https://github.com/microsoft/satclip.git" @@ -37,3 +64,7 @@ Issues = "https://github.com/microsoft/satclip/issues" [tool.setuptools.packages.find] include = ["satclip", "satclip.*"] + +[tool.setuptools.package-data] +# Ship the default training config so `python -m satclip.main` works after install. +satclip = ["configs/*.yaml"] diff --git a/satclip/__init__.py b/satclip/__init__.py index cdbeb53..d563d64 100644 --- a/satclip/__init__.py +++ b/satclip/__init__.py @@ -1,5 +1,11 @@ -from . import * -from .main import * -from .model import * -from .loss import * -from .location_encoder import * \ No newline at end of file +"""SatCLIP: a global, general-purpose geographic location encoder. + +See https://github.com/microsoft/satclip for details. +""" + +from .load import get_satclip +from .location_encoder import LocationEncoder +from .loss import SatCLIPLoss +from .model import SatCLIP + +__all__ = ["SatCLIP", "LocationEncoder", "SatCLIPLoss", "get_satclip"] diff --git a/satclip/main.py b/satclip/main.py index a5134ab..88e6a24 100644 --- a/satclip/main.py +++ b/satclip/main.py @@ -11,6 +11,10 @@ torch.set_float32_matmul_precision('high') +# Default training config shipped with the package; resolved relative to this +# file so training works regardless of the current working directory. +DEFAULT_CONFIG = Path(__file__).resolve().parent / "configs" / "default.yaml" + class SatCLIPLightningModule(lightning.pytorch.LightningModule): def __init__( self, @@ -112,7 +116,9 @@ def add_arguments_to_parser(self, parser): parser.add_argument("--watchmodel", action="store_true") -def cli_main(default_config_filename="./configs/default.yaml"): +def cli_main(default_config_filename=None): + if default_config_filename is None: + default_config_filename = str(DEFAULT_CONFIG) save_config_fn = default_config_filename.replace(".yaml", "-latest.yaml") # modify configs/default.yaml for learning rate etc. cli = MyLightningCLI( @@ -140,7 +146,7 @@ def cli_main(default_config_filename="./configs/default.yaml"): # Create folder to log configs # NOTE: Lightning does not handle config paths with subfolders - dirname_cfg = Path(default_config_filename).parent + dirname_cfg = Path(default_config_filename).parent.name dir_log_cfg = Path(cli.trainer.log_dir) / dirname_cfg dir_log_cfg.mkdir(parents=True, exist_ok=True) @@ -151,12 +157,10 @@ def cli_main(default_config_filename="./configs/default.yaml"): if __name__ == "__main__": - config_fn = "./configs/default.yaml" - #A100 go vroom vroom 🚗💨 - if torch.cuda.get_device_name(device=0)=='NVIDIA A100 80GB PCIe': + if torch.cuda.is_available() and torch.cuda.get_device_name(device=0) == 'NVIDIA A100 80GB PCIe': torch.backends.cuda.matmul.allow_tf32 = True print('Superfastmode! 🚀') else: torch.backends.cuda.matmul.allow_tf32 = False - cli_main(config_fn) \ No newline at end of file + cli_main() \ No newline at end of file diff --git a/tests/test_smoke.py b/tests/test_smoke.py new file mode 100644 index 0000000..ce31b89 --- /dev/null +++ b/tests/test_smoke.py @@ -0,0 +1,71 @@ +"""Smoke tests for the installed `satclip` package. + +These verify that the package imports cleanly (i.e. that intra-package imports +resolve as a proper package) and that a small SatCLIP model can run a forward +pass end to end. +""" + +import torch + +import satclip +from satclip.model import SatCLIP + + +def test_public_api(): + for name in ("SatCLIP", "LocationEncoder", "SatCLIPLoss", "get_satclip"): + assert hasattr(satclip, name), f"satclip.{name} should be importable" + + +def test_satclip_forward(): + model = SatCLIP( + embed_dim=128, + image_resolution=224, + in_channels=13, + vision_layers=2, + vision_width=64, + vision_patch_size=32, + le_type="sphericalharmonics", + pe_type="siren", + legendre_polys=10, + frequency_num=16, + max_radius=360, + min_radius=1, + harmonics_calculation="analytic", + ) + model.eval() + + img_batch = torch.randn(2, 13, 224, 224) + loc_batch = torch.randn(2, 2) + + with torch.no_grad(): + logits_per_image, logits_per_coord = model(img_batch, loc_batch) + + assert logits_per_image.shape == (2, 2) + assert logits_per_coord.shape == (2, 2) + + +def test_location_encoder_only(): + """The location encoder should embed coordinates on its own.""" + model = SatCLIP( + embed_dim=128, + image_resolution=224, + in_channels=13, + vision_layers=2, + vision_width=64, + vision_patch_size=32, + le_type="sphericalharmonics", + pe_type="siren", + legendre_polys=10, + frequency_num=16, + max_radius=360, + min_radius=1, + harmonics_calculation="analytic", + ) + model.eval() + + coords = torch.randn(4, 2) + with torch.no_grad(): + emb = model.encode_location(coords) + + assert emb.shape[0] == 4 + assert emb.shape[1] == 128