diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..fed552b --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,51 @@ +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: + # Oldest supported and newest Python. + python-version: ["3.9", "3.13"] + + 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: Build sdist and wheel + run: | + pip install build + python -m build + + - name: Install wheel + # Install the built wheel (not an editable install) so packaging mistakes, + # e.g. missing package data, are caught. + run: pip install "$(ls dist/*.whl)[dev,examples]" + + - name: Run tests against the installed package + # Run from outside the repository so `import satclip` resolves to the + # installed wheel rather than the source tree. + run: | + cp -r tests "$RUNNER_TEMP/tests" + cd "$RUNNER_TEMP" + python -c "import satclip; print('satclip public API:', satclip.__all__)" + satclip-train --help > /dev/null + pytest -q tests diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e863259 --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +# SatCLIP files +satclip_logs/ + +# Python +__pycache__/ +*.egg-info/ +build/ +dist/ +.pytest_cache/ diff --git a/README.md b/README.md index 31ac775..810ef8f 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 (Python 3.9+). 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, run the `satclip-train` command (installed with the package), pointing it at this dataset directory: ```bash -cd satclip -python main.py +satclip-train --data.data_dir /path/to/s2-100k ``` +Training uses the default config in [`satclip/configs/default.yaml`](satclip/configs/default.yaml). Any option can be overridden on the command line (see `satclip-train --help`), or you can pass your own config, e.g. a modified copy of the default: `satclip-train --config path/to/config.yaml`. Logs, checkpoints and the resolved config (`default-latest.yaml`) are written to `satclip_logs/` in the current directory. `python -m satclip.main` is equivalent. ### 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..7aed25f 100644 --- a/notebooks/A01_Simple_SatCLIP_Usage.ipynb +++ b/notebooks/A01_Simple_SatCLIP_Usage.ipynb @@ -19,7 +19,7 @@ "source": [ "## A01 - Simple usage of pretrained SatCLIP embeddings\n", "\n", - "To obtained pretrained **SatCLIP** embeddings, first install the repository." + "To obtain pretrained **SatCLIP** embeddings, first install the `satclip` package (this also installs all required dependencies)." ], "metadata": { "id": "ngz8zz9Gvbxh" @@ -51,68 +51,7 @@ } ], "source": [ - "!rm -r sample_data .config # Empty current directory\n", - "!git clone https://github.com/microsoft/satclip.git . # Clone SatCLIP repository" - ] - }, - { - "cell_type": "markdown", - "source": [ - "Now install the required Python packages." - ], - "metadata": { - "id": "hOEZnNt1v2Bl" - } - }, - { - "cell_type": "code", - "source": [ - "!pip install lightning --quiet\n", - "!pip install rasterio --quiet\n", - "!pip install torchgeo --quiet" - ], - "metadata": { - "colab": { - "base_uri": "https://localhost:8080/" - }, - "id": "Q72Ypu0Cr3Sc", - "outputId": "d54f24b8-26b1-4f86-ffc8-23da0cc63710" - }, - "execution_count": null, - "outputs": [ - { - "output_type": "stream", - "name": "stdout", - "text": [ - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.0/2.0 MB\u001b[0m \u001b[31m13.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m806.1/806.1 kB\u001b[0m \u001b[31m40.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m776.9/776.9 kB\u001b[0m \u001b[31m42.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m20.6/20.6 MB\u001b[0m \u001b[31m27.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m378.5/378.5 kB\u001b[0m \u001b[31m5.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m44.6/44.6 kB\u001b[0m \u001b[31m5.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m705.7/705.7 kB\u001b[0m \u001b[31m37.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m670.7/670.7 kB\u001b[0m \u001b[31m7.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m488.6/488.6 kB\u001b[0m \u001b[31m41.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m106.7/106.7 kB\u001b[0m \u001b[31m11.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.2/2.2 MB\u001b[0m \u001b[31m66.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m154.5/154.5 kB\u001b[0m \u001b[31m15.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m137.6/137.6 kB\u001b[0m \u001b[31m13.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m92.6/92.6 MB\u001b[0m \u001b[31m7.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m189.7/189.7 kB\u001b[0m \u001b[31m14.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m79.5/79.5 kB\u001b[0m \u001b[31m8.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m101.7/101.7 kB\u001b[0m \u001b[31m6.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.8/58.8 kB\u001b[0m \u001b[31m5.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[?25h Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n", - " Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.2/2.2 MB\u001b[0m \u001b[31m21.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m117.0/117.0 kB\u001b[0m \u001b[31m9.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[?25h Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m594.2/594.2 kB\u001b[0m \u001b[31m44.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[?25h Building wheel for efficientnet-pytorch (setup.py) ... \u001b[?25l\u001b[?25hdone\n", - " Building wheel for pretrainedmodels (setup.py) ... \u001b[?25l\u001b[?25hdone\n", - " Building wheel for antlr4-python3-runtime (setup.py) ... \u001b[?25l\u001b[?25hdone\n" - ] - } + "!pip install git+https://github.com/microsoft/satclip.git" ] }, { @@ -169,11 +108,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..087b185 100644 --- a/notebooks/A02_SatCLIP_Hugging_Face_Usage.ipynb +++ b/notebooks/A02_SatCLIP_Hugging_Face_Usage.ipynb @@ -365,7 +365,7 @@ "source": [ "## A02 - Simple usage of pretrained SatCLIP embeddings using Hugging Face\n", "\n", - "To obtained pretrained **SatCLIP** embeddings, first install the repository." + "To obtain pretrained **SatCLIP** embeddings, first install the `satclip` package (this also installs all required dependencies)." ], "metadata": { "id": "ngz8zz9Gvbxh" @@ -397,68 +397,7 @@ } ], "source": [ - "!rm -r sample_data .config # Empty current directory\n", - "!git clone https://github.com/microsoft/satclip.git . # Clone SatCLIP repository" - ] - }, - { - "cell_type": "markdown", - "source": [ - "Now install the required Python packages." - ], - "metadata": { - "id": "hOEZnNt1v2Bl" - } - }, - { - "cell_type": "code", - "source": [ - "!pip install lightning --quiet\n", - "!pip install rasterio --quiet\n", - "!pip install torchgeo --quiet" - ], - "metadata": { - "colab": { - "base_uri": "https://localhost:8080/" - }, - "id": "Q72Ypu0Cr3Sc", - "outputId": "7f10688e-5f74-4e8d-ce5b-9c2fa2219869" - }, - "execution_count": 2, - "outputs": [ - { - "output_type": "stream", - "name": "stdout", - "text": [ - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.0/2.0 MB\u001b[0m \u001b[31m9.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m840.2/840.2 kB\u001b[0m \u001b[31m15.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m777.7/777.7 kB\u001b[0m \u001b[31m14.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m20.6/20.6 MB\u001b[0m \u001b[31m34.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m378.5/378.5 kB\u001b[0m \u001b[31m5.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m44.6/44.6 kB\u001b[0m \u001b[31m3.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m756.0/756.0 kB\u001b[0m \u001b[31m26.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m684.1/684.1 kB\u001b[0m \u001b[31m41.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m535.2/535.2 kB\u001b[0m \u001b[31m2.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m106.7/106.7 kB\u001b[0m \u001b[31m10.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.2/2.2 MB\u001b[0m \u001b[31m53.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m154.5/154.5 kB\u001b[0m \u001b[31m16.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m137.6/137.6 kB\u001b[0m \u001b[31m13.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m105.0/105.0 MB\u001b[0m \u001b[31m7.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m189.8/189.8 kB\u001b[0m \u001b[31m14.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m79.5/79.5 kB\u001b[0m \u001b[31m7.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m101.7/101.7 kB\u001b[0m \u001b[31m10.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.8/58.8 kB\u001b[0m \u001b[31m5.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[?25h Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n", - " Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.2/2.2 MB\u001b[0m \u001b[31m67.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m117.0/117.0 kB\u001b[0m \u001b[31m12.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[?25h Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m594.2/594.2 kB\u001b[0m \u001b[31m41.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[?25h Building wheel for efficientnet-pytorch (setup.py) ... \u001b[?25l\u001b[?25hdone\n", - " Building wheel for pretrainedmodels (setup.py) ... \u001b[?25l\u001b[?25hdone\n", - " Building wheel for antlr4-python3-runtime (setup.py) ... \u001b[?25l\u001b[?25hdone\n" - ] - } + "!pip install git+https://github.com/microsoft/satclip.git" ] }, { @@ -473,11 +412,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 +434,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..a766d8a 100644 --- a/notebooks/B01_Example_Air_Temperature_Prediction.ipynb +++ b/notebooks/B01_Example_Air_Temperature_Prediction.ipynb @@ -55,16 +55,12 @@ } ], "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" ] }, { "cell_type": "code", "source": [ - "!pip install lightning --quiet\n", - "!pip install rasterio --quiet\n", - "!pip install torchgeo --quiet\n", "!pip install basemap --quiet" ], "metadata": { @@ -168,11 +164,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..eb65c90 100644 --- a/notebooks/B02_Example_Image_Localization.ipynb +++ b/notebooks/B02_Example_Image_Localization.ipynb @@ -399,16 +399,12 @@ } ], "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" ] }, { "cell_type": "code", "source": [ - "!pip install lightning --quiet\n", - "!pip install rasterio --quiet\n", - "!pip install torchgeo --quiet\n", "!pip install basemap --quiet" ], "metadata": { @@ -485,10 +481,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 new file mode 100644 index 0000000..d367df0 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,71 @@ +[build-system] +requires = ["setuptools >= 77.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 = "MIT" +license-files = ["LICENSE"] +requires-python = ">=3.9" +classifiers = [ + "Programming Language :: Python :: 3", + "Operating System :: OS Independent", +] + +# 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", + "einops", + "huggingface_hub", + "jsonargparse[signatures] >=4.27.7", # LightningCLI (satclip-train) + "lightning >=2.2,<3", + "matplotlib", + "numpy", + "pandas", + "rasterio >=1.3.10", + "sympy", + "tensorboard", # default training logger (configs/default.yaml) + "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", +] + +[project.scripts] +satclip-train = "satclip.main:cli_main" + +[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.*"] + +[tool.setuptools.package-data] +# Ship the default training config so `python -m satclip.main` works after install. +satclip = ["configs/*.yaml"] diff --git a/satclip/.gitignore b/satclip/.gitignore deleted file mode 100644 index 52804a2..0000000 --- a/satclip/.gitignore +++ /dev/null @@ -1,5 +0,0 @@ -# SatCLIP files -satclip_logs/ - -# Python -__pycache__/ \ No newline at end of file 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/load.py b/satclip/load.py index ac975de..581166e 100644 --- a/satclip/load.py +++ b/satclip/load.py @@ -1,11 +1,15 @@ -from main import * +import inspect + +from .main import * def get_satclip(ckpt_path, device, return_all=False): ckpt = torch.load(ckpt_path,map_location=device) - ckpt['hyper_parameters'].pop('eval_downstream') - ckpt['hyper_parameters'].pop('air_temp_data_path') - ckpt['hyper_parameters'].pop('election_data_path') - lightning_model = SatCLIPLightningModule(**ckpt['hyper_parameters']).to(device) + # Checkpoints can carry hyperparameters SatCLIPLightningModule does not accept: + # released ones have legacy keys (e.g. eval_downstream), and ones trained via + # LightningCLI have '_instantiator'. Keep only the constructor's arguments. + accepted = inspect.signature(SatCLIPLightningModule.__init__).parameters + hparams = {k: v for k, v in ckpt['hyper_parameters'].items() if k in accepted} + lightning_model = SatCLIPLightningModule(**hparams).to(device) lightning_model.load_state_dict(ckpt['state_dict']) lightning_model.eval() @@ -15,4 +19,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..1f52a0a 100644 --- a/satclip/main.py +++ b/satclip/main.py @@ -3,13 +3,18 @@ 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') +# 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, @@ -111,8 +116,20 @@ def add_arguments_to_parser(self, parser): parser.add_argument("--watchmodel", action="store_true") -def cli_main(default_config_filename="./configs/default.yaml"): - save_config_fn = default_config_filename.replace(".yaml", "-latest.yaml") +def configure_tf32(): + #A100 go vroom vroom 🚗💨 + 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! 🚀') + + +def cli_main(default_config_filename=None, args=None): + if default_config_filename is None: + default_config_filename = str(DEFAULT_CONFIG) + # Saved next to the run's logs (Lightning resolves it relative to trainer.log_dir), + # never next to the source config, which may live inside the installed package. + save_config_fn = Path(default_config_filename).stem + "-latest.yaml" + configure_tf32() # modify configs/default.yaml for learning rate etc. cli = MyLightningCLI( model_class=SatCLIPLightningModule, @@ -128,6 +145,7 @@ def cli_main(default_config_filename="./configs/default.yaml"): parser_kwargs={"default_config_files": [default_config_filename]}, seed_everything_default=0, run=False, + args=args, ) ts = datetime.now().strftime("%Y-%m-%d_%H:%M:%S") @@ -137,12 +155,6 @@ def cli_main(default_config_filename="./configs/default.yaml"): # this seems to be necessary to force logging of datamodule hyperparams cli.trainer.logger.log_hyperparams(cli.datamodule.hparams) - # Create folder to log configs - # NOTE: Lightning does not handle config paths with subfolders - dirname_cfg = Path(default_config_filename).parent - dir_log_cfg = Path(cli.trainer.log_dir) / dirname_cfg - dir_log_cfg.mkdir(parents=True, exist_ok=True) - cli.trainer.fit( model=cli.model, datamodule=cli.datamodule, @@ -150,12 +162,4 @@ 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': - 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() 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) 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 diff --git a/tests/test_train.py b/tests/test_train.py new file mode 100644 index 0000000..8f7f2c5 --- /dev/null +++ b/tests/test_train.py @@ -0,0 +1,98 @@ +"""End-to-end training smoke test. + +Trains for a couple of steps on a tiny synthetic S2-100K-style dataset through +the same entry point as `satclip-train`, then reloads the checkpoint with +`get_satclip`. This exercises the packaged default config, the data module and +the checkpoint round trip. +""" + +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +import rasterio +import torch + +import satclip +from satclip.load import get_satclip +from satclip.main import SatCLIPLightningModule, cli_main + +N_TILES = 12 +N_BANDS = 12 # S2-100K tiles omit B10; the training transform re-inserts it +TILE_SIZE = 256 + + +@pytest.fixture +def s2_dataset(tmp_path): + root = tmp_path / "s2" + (root / "images").mkdir(parents=True) + rng = np.random.default_rng(0) + # The dataset integrity check expects the first and last S2-100K patch names. + ids = [0, *range(1, N_TILES - 1), 99999] + rows = [] + for i in ids: + fn = f"patch_{i}.tif" + with rasterio.open( + root / "images" / fn, "w", driver="GTiff", + height=TILE_SIZE, width=TILE_SIZE, count=N_BANDS, dtype="uint16", + ) as f: + f.write(rng.integers(0, 10000, (N_BANDS, TILE_SIZE, TILE_SIZE), dtype=np.uint16)) + rows.append({"fn": fn, "lon": rng.uniform(-180, 180), "lat": rng.uniform(-90, 90)}) + pd.DataFrame(rows).to_csv(root / "index.csv", index=False) + return root + + +@pytest.mark.filterwarnings("ignore") +def test_train_and_reload(s2_dataset, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + package_configs = Path(satclip.__file__).parent / "configs" + configs_before = sorted(p.name for p in package_configs.iterdir()) + + cli_main(args=[ + f"--data.data_dir={s2_dataset}", + "--data.batch_size=4", + "--data.num_workers=0", + "--trainer.accelerator=cpu", + "--trainer.max_epochs=1", + "--trainer.limit_train_batches=2", + "--trainer.limit_val_batches=1", + "--trainer.accumulate_grad_batches=1", + "--trainer.enable_progress_bar=false", + ]) + + log_dir = tmp_path / "satclip_logs" / "satclip" / "satclip" + assert (log_dir / "default-latest.yaml").is_file() + # Training must not write into the installed package. + assert sorted(p.name for p in package_configs.iterdir()) == configs_before + + ckpt_path = log_dir / "checkpoints" / "last.ckpt" + assert ckpt_path.is_file() + + location_encoder = get_satclip(ckpt_path, device="cpu") + with torch.no_grad(): + emb = location_encoder(torch.tensor([[-74.0, 40.7], [13.4, 52.5]]).double()) + assert emb.shape == (2, 256) + assert torch.isfinite(emb).all() + + +def test_get_satclip_accepts_released_checkpoint_hparams(tmp_path): + """Released checkpoints carry extra legacy hyperparameters that must be dropped.""" + module = SatCLIPLightningModule( + embed_dim=32, image_resolution=64, vision_layers=2, vision_width=64, + vision_patch_size=32, in_channels=13, legendre_polys=4, + ) + ckpt = { + "state_dict": module.state_dict(), + "hyper_parameters": { + **dict(module.hparams), + "eval_downstream": False, + "air_temp_data_path": None, + "election_data_path": None, + }, + } + ckpt_path = tmp_path / "released.ckpt" + torch.save(ckpt, ckpt_path) + + model = get_satclip(ckpt_path, device="cpu", return_all=True) + assert isinstance(model, satclip.SatCLIP)