Reusable helpers for interacting with the tytonai platform: S3 downloads, web map tiling, model fetch, and annotation-mask rollup. Built to be installed straight from GitHub across projects.
- Web map import — cut local GeoTIFF tiles from an S3 web map (any band layout). ✅
- Manifest import — download tiles listed in a manifest (image + mask
.npz). ✅ - Model fetch — download trained weights, or build a fresh model, from a config. ✅
- Mask rollup — remap annotation-mask class ids to parent categories. ✅
- Annotation realign — re-tile misaligned annotation masks onto a grid. ✅
A full, runnable line-by-line walkthrough of every feature lives at
examples/tutorial.py — it states the input files each section
needs and is the fastest way to smoke-test the whole package after install.
Install from GitHub. The minimal consumer pyproject.toml is kept in the repo at
examples/pyproject.toml — copy it into your project and trim
the extras to the features you use. It is updated whenever a new facility is added.
# minimal consumer pyproject.toml (see examples/pyproject.toml)
[project]
name = "your-project"
version = "0.1.0"
requires-python = ">=3.14"
dependencies = ["tytonai_utils[webmap,s3,model]"] # trim extras as needed
[tool.uv.sources]
tytonai_utils = { git = "https://github.com/clement-tyton/tytonai_utils.git" }uv sync
# re-pull after the package is updated:
uv lock --upgrade-package tytonai_utils && uv sync| Extra | Pulls in | For |
|---|---|---|
webmap |
geopandas, rasterio, matplotlib, rio_tiler, bbox_to_tile_grid | Feature 1 |
s3 |
boto3 | Features 2–4 (S3 API) |
model |
torch, segmentation-models-pytorch, timm | Feature 3 (build/load model) |
viz |
matplotlib | Visualization helpers (image/mask QA plots) |
all |
everything above | — |
All credentials/config come from a .env file (loaded with python-dotenv). Copy
.env.example and fill it in. Never commit .env — it is gitignored.
The AWS_* keys are exactly GDAL's /vsis3 config variables, so for the web map feature
load_dotenv() is the only auth step needed — GDAL reads them automatically. Key fields:
| Key | Purpose |
|---|---|
AWS_S3_ENDPOINT |
S3-compatible host, e.g. s3.tytonai.com |
AWS_HTTPS, AWS_VIRTUAL_HOSTING |
YES / FALSE (path-style addressing) |
AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_SESSION_TOKEN |
credentials |
S3_FILE_BUCKET |
default data bucket |
A web map is a Cloud-Optimized GeoTIFF on S3. Tiles are never pre-downloaded: GDAL
opens the raster via /vsis3 and does ranged GETs, reading only each tile's window.
from dotenv import load_dotenv
from tytonai_utils.webmap import download_webmap_from_shp
load_dotenv()
# one call: vector area + S3 link -> tiles
written = download_webmap_from_shp(
"study_area.fgb",
"s3://bucket/id/RED_GREEN_BLUE_NIR_ALPHA_webmap.tif",
"tiles_out",
res=0.1, patch=512,
bands=[1, 2, 3], # RGB; omit for all bands except alpha
)download_webmap_from_shp(shp_path, webmap, out_dir, res, patch, bands=None, workers=8, skip_empty=True) -> list[str]
High-level one-call entry: build the tile grid from a vector file, then download each tile from the web map. Use this unless you need to inspect/modify the grid first.
| Param | Type | Description |
|---|---|---|
shp_path |
str | Path |
Vector area (.shp / .fgb / .geojson — anything geopandas reads) |
webmap |
str |
Web map link (s3://, https://, /vsis3/…) — auto-normalized |
out_dir |
str | Path |
Output folder |
res, patch |
float, int |
The web map's native resolution (m/px — from the tytonai app, beside the S3 link) and tile size (px) |
bands |
list[int] | None |
Bands to write (see download_grid) |
workers, skip_empty |
int, bool |
Passed through to download_grid |
Returns the list of written tile filenames. Internally calls build_grid +
download_grid — use those directly for finer control (e.g. plotting the grid, slicing
to a subset before downloading).
Normalize a web map link into a GDAL-openable path for ranged reads. Accepts and converts:
s3://bucket/key→/vsis3/bucket/keyhttp(s)://…→/vsicurl/…/vsis3/…,/vsicurl/…, or a local path → passed through unchanged
You rarely call this directly — download_grid applies it for you, so you can paste the
s3:// link straight from the app.
| Param | Type | Description |
|---|---|---|
link |
str |
Web map link in any of the forms above |
Returns the GDAL path string.
Read a .fgb (FlatGeobuf vector area), take its bounding box, and build a uniform
tile grid in pixel space anchored at the bbox top-left.
| Param | Type | Description |
|---|---|---|
fgb_path |
str | Path |
FlatGeobuf file defining the region to tile |
res |
float |
The web map's native resolution (metres/pixel) — read it in the tytonai app, in the same place you copy the S3 link |
patch |
int |
Tile size in pixels (square) |
Returns (grid, study_area) — both GeoDataFrames. grid has one row per tile;
study_area is the original vector area (used for plotting / CRS).
Download every tile of grid from webmap into out_dir as georeferenced GeoTIFFs.
Reads run in parallel (I/O-bound; each worker thread opens its own rasterio handle).
| Param | Type | Description |
|---|---|---|
grid |
GeoDataFrame |
Tile grid from build_grid (slice it to limit, e.g. grid.iloc[:10]) |
webmap |
str |
Web map link (s3://, https://, /vsis3/…) — auto-normalized |
out_dir |
str | Path |
Output folder (created if missing) |
bands |
list[int] | None |
1-based bands to write. [1,2,3]=RGB, [1,2,3,4]=RGB+NIR, [1]=mask. None = every band except alpha |
workers |
int |
Parallel download threads (default 8) |
skip_empty |
bool |
Skip tiles with no coverage (alpha→nodata→nonzero) |
Returns the list of written filenames, named by grid index (tile_<cell index>.tif), so a
subset via grid.iloc[...] keeps matching names. Empty tiles are omitted. Output preserves the
source dtype, nodata, CRS, and per-band colour interpretation — so masks stay clean integer
rasters.
Save a PNG of the tile grid (blue) over the study-area outline (red) — a quick sanity check of coverage and tile count before downloading. The cell size in the title is read from the grid geometry.
| Param | Type | Description |
|---|---|---|
grid, study_area |
GeoDataFrame |
Outputs of build_grid |
name |
str |
Title label |
out_png |
str | Path |
Output PNG path |
Coarse mosaic of downloaded tiles: each .tif read downsampled and placed at its real
geo-extent. Handles RGB (3+ bands) and single-band (grayscale) tiles.
| Param | Type | Description |
|---|---|---|
tiles_dir |
str | Path |
Folder of downloaded .tif tiles |
downscale |
int |
Read 1/downscale resolution for speed |
ax |
matplotlib Axes | None |
Draw onto an existing axes (compare areas side by side) |
out_png |
str | Path | None |
Save the figure if given |
Returns the matplotlib Axes. Raises FileNotFoundError if no tiles are present.
Read a dataset.json manifest (a list of tile dicts) and download every imagery + mask
.npz it references from S3 into a chosen out_dir. Uses boto3 (the s3 extra) — no
aws CLI needed. Cache-aware: files already on disk are skipped unless force=True.
from dotenv import load_dotenv
from tytonai_utils.manifest import download_annotations_from_dataset_manifest
load_dotenv()
download_annotations_from_dataset_manifest("monrovia/dataset.json", out_dir="monrovia/annotations")Load the dataset.json tile list (JSON array of tile dicts with imagery_file,
mask_file, geotransform, srid, class_counts, …).
| Param | Type | Description |
|---|---|---|
manifest_path |
str | Path |
Path to the dataset.json file |
Download s3://bucket/key to dest. Returns True if a download happened, False if
the file was already cached on disk.
| Param | Type | Description |
|---|---|---|
s3 |
boto3 client | From make_s3_client() |
key |
str |
S3 object key (the manifest's imagery_file / mask_file) |
dest |
Path |
Local destination (parent dirs created) |
bucket |
str |
Source bucket |
force |
bool |
Re-download even if dest exists |
download_annotations_from_dataset_manifest(manifest_path, out_dir, bucket=None, force=False, workers=8, s3=None) -> Path
Download every imagery + mask NPZ referenced by the manifest into out_dir, in parallel.
| Param | Type | Description |
|---|---|---|
manifest_path |
str | Path |
Path to dataset.json |
out_dir |
str | Path |
Output folder (created if missing) |
bucket |
str | None |
Source bucket; defaults to $S3_FILE_BUCKET |
force |
bool |
Re-download cached files |
workers |
int |
Parallel download threads (default 8) |
s3 |
boto3 client | None | Reuse a client; defaults to a fresh make_s3_client() |
Returns out_dir. Prints a summary (downloaded vs cached). Requires load_dotenv()
beforehand for the AWS_* creds + endpoint.
A model config (JSON) describes a segmentation_models_pytorch (smp) model — model_type,
encoder_type, encoder_weights, bands (→ in_channels), class_list (→ #classes), and
epoch_file_key (the s3://…/...pth trained weights). This feature downloads the weights
and/or instantiates the model — fresh, loaded, or reshaped for transfer learning.
Heavy deps are lazy-imported: the download function needs only the s3 extra; the build/load
functions need the model extra (torch + segmentation-models-pytorch).
from dotenv import load_dotenv
from tytonai_utils.model import (
build_model_from_config,
download_model_weights_from_config,
load_model_with_fresh_head_from_config,
load_trained_model_from_config,
)
load_dotenv()
# fresh model, no checkpoint: ImageNet encoder + random head
fresh = build_model_from_config("model_config.json")
# download the trained weights, then load — choose ONE of the two load modes:
weights = download_model_weights_from_config("model_config.json", "models/")
# (a) exactly as trained — full weights incl. head (same classes as the checkpoint)
trained = load_trained_model_from_config("model_config.json", weights)
# (b) finetune — reuse encoder+decoder, fresh random head for a new 3-class set
finetune = load_model_with_fresh_head_from_config(
"model_config.json", weights, num_classes=3, freeze_encoder=True
)Which load function? Only the segmentation head depends on the class count; the encoder and decoder always transfer. Use
load_trained_model_from_configwhen you want the model as-is (it errors if the checkpoint's class count differs). Useload_model_with_fresh_head_from_configwhen the class set changes — it reuses encoder+decoder and leaves a fresh random head.
Load a model config JSON.
Download the trained .pth referenced by the config into out_dir. The S3 link and its
bucket come from the config's epoch_file_key (not $S3_FILE_BUCKET). Cache-aware.
Returns the local weights path. (s3 extra.)
Build the smp architecture from the config. Default = ImageNet-pretrained encoder + randomly
initialised decoder/head. pretrained_encoder=False → fully random. num_classes overrides
the head size (else len(class_list)). (model extra.)
| Param | Type | Description |
|---|---|---|
config_path |
str | Path |
Model config JSON |
pretrained_encoder |
bool |
Load ImageNet encoder weights (head always random) |
num_classes |
int | None |
Override the head class count |
load_trained_model_from_config(config_path, weights_path, freeze_encoder=False, map_location="cpu") -> nn.Module
Load the model exactly as trained — full weights including the segmentation head.
Loads strictly. Raises if the checkpoint's class count differs from the config's
class_list (use the fresh-head function for a different class set). Handles wrapped
checkpoints ({"model": …} from Fabric/training loops) and plain state-dicts.
| Param | Type | Description |
|---|---|---|
config_path |
str | Path |
Model config JSON |
weights_path |
str | Path |
Local .pth checkpoint (from the download fn) |
freeze_encoder |
bool |
Set requires_grad=False on the encoder |
map_location |
str |
Device for torch.load (e.g. "cpu", "cuda") |
load_model_with_fresh_head_from_config(config_path, weights_path, num_classes=None, freeze_encoder=False, map_location="cpu") -> nn.Module
Finetune setup: reuse the checkpoint's encoder + decoder, with a fresh random head
sized to num_classes (defaults to len(class_list)). Head weights in the checkpoint are
dropped; everything else must load or it raises (so a key mismatch can't silently
produce a random model).
| Param | Type | Description |
|---|---|---|
config_path |
str | Path |
Model config JSON |
weights_path |
str | Path |
Local .pth checkpoint |
num_classes |
int | None |
New head size (else len(class_list)) |
freeze_encoder |
bool |
Set requires_grad=False on the encoder |
map_location |
str |
Device for torch.load |
Remap annotation-mask class ids to parent categories. The grouping (which source classes
roll into which parent) is the source of truth; a remapping is that grouping + a number
scheme ({parent_name: id}). Ids not in the grouping go to a nodata value (default 0).
Pure numpy + stdlib — no extra needed.
Two R&D schemes ship in the module:
- 7-class — Ground, Shrub, Tree, Herb, Sedge, Tussock, Hummock (
RND_REMAP_7CLASS). - 6-class — same, but Tussock + Hummock + generic Grass fold into Grass (
RND_REMAP_6CLASS).
In both, unmapped classes (Biotic, Not Erosion, Erosion — plus Grass in the 7-class) → nodata.
from tytonai_utils.rollup import RND_REMAP_7CLASS, rollup_mask, rollup_annotations
remapped = rollup_mask(mask, RND_REMAP_7CLASS) # one numpy mask
rollup_annotations("annotations/", "dataset.json", # a whole folder
RND_REMAP_7CLASS, out_dir="annotations_rnd7/")| Name | What |
|---|---|
CLASS_NAMES |
source {id: name} (the org class list) |
ROLLUP_GROUPS_7CLASS / ROLLUP_GROUPS_6CLASS |
{parent_name: [source_ids]} groupings |
RND_TARGET_IDS_7CLASS / _6CLASS |
{parent_name: target_id} number schemes |
RND_REMAP_7CLASS / _6CLASS |
materialized {source_id: target_id} |
RND_NAMES_7CLASS / _6CLASS |
target {id: name} for the remapped masks |
NODATA |
value for unmapped ids (default 0) |
Flatten a grouping ({name: [source_ids]}) + a number scheme ({name: id}) into a
{source_id: target_id} dict. Use this to define a new remapping that reuses an existing
grouping with different numbers.
The remapped class list as {target_id: parent_name}.
Vectorized remap of one integer mask. Ids not in remapping → nodata (pass nodata=None
to keep them unchanged). Returns a new array of the same dtype; input untouched.
rollup_annotations(annotations_dir, manifest, remapping, out_dir, mask_key=None, nodata=0) -> list[Path]
Roll up every mask .npz referenced by a manifest, writing remapped masks to out_dir.
Imagery files untouched; mask array key auto-detected (largest) unless mask_key given.
| Param | Type | Description |
|---|---|---|
annotations_dir |
str | Path |
Folder of downloaded mask .npz |
manifest |
list[dict] | str | Path |
Tile list or path to dataset.json |
remapping |
dict[int,int] |
{source_id: target_id} (e.g. RND_REMAP_7CLASS) |
out_dir |
str | Path |
Output folder for remapped masks |
mask_key |
str | None |
NPZ mask key; auto-detected if None |
nodata |
int | None |
Value for unmapped ids (default 0; None = keep) |
Annotation tiles from a manifest are georeferenced (geotransform + srid) but rarely tiled
on the same grid as your web map. This merges every annotation mask into one mosaic and
cuts it along your grid cells — masks aligned 1:1 with the imagery download_grid writes
for the same grid (output tile_NNNNN.tif pairs with imagery tile_NNNNN.tif).
Uses rasterio (the webmap extra) + numpy. Match the grid res to the annotation resolution
(resolution_values in the manifest) for a clean re-tile.
from tytonai_utils.webmap import build_grid, download_grid
from tytonai_utils.align import realign_annotations_to_grid
# ONE grid drives BOTH imagery and masks so the .tif tiles pair by filename.
grid, _ = build_grid("study_area.fgb", res=0.0206348504972787, patch=512) # res = annotation res
download_grid(grid, webmap, "tiles_aligned/", bands=[1, 2, 3]) # imagery on the grid
realign_annotations_to_grid(grid, "annotations/", "dataset.json", "masks_aligned/") # masks on SAME grid
# tiles_aligned/tile_NNNNN.tif <-> masks_aligned/tile_NNNNN.tifPairing rule. Imagery (
download_grid) and masks (realign_annotations_to_grid) must be built from the same grid object (samefgb/res/patch) — both name tilestile_NNNNN.tifby grid index (the cell's position in the GeoDataFrame), so a different grid (e.g. a differentres) yields non-matching filenames andplot_image_mask_tilesfinds no pairs.Cheap imagery for annotated cells only. Because tiles are named by grid index, a subset keeps its names: realign first to learn which cells have annotations, then download imagery for just those cells —
aligned = realign_annotations_to_grid(grid, "annotations/", "dataset.json", "masks_aligned/") cells = sorted(int(n.removeprefix("tile_").removesuffix(".tif")) for n in aligned) download_grid(grid.iloc[cells], webmap, "tiles_aligned/", bands=[1, 2, 3]) # only those cells
grid.iloc[cells]preserves the index labels, sotiles_aligned/tile_NNNNN.tifmatchesmasks_aligned/tile_NNNNN.tif.
realign_annotations_to_grid(grid, annotations_dir, manifest, out_dir, mask_key=None, nodata=0, skip_empty=True, overlapping="first") -> list[str]
Mosaic the annotation masks and cut them along grid (a GeoDataFrame from build_grid or any
tiling). Writes grid-aligned GeoTIFF masks; skips cells with no annotation coverage. The grid
is reprojected to the annotation CRS if they differ. Shows tqdm progress (load → mosaic → cut).
| Param | Type | Description |
|---|---|---|
grid |
GeoDataFrame |
Target tiling (from build_grid) |
annotations_dir |
str | Path |
Folder of downloaded mask .npz |
manifest |
list[dict] | str | Path |
Tile list or path to dataset.json (provides geotransform/srid) |
out_dir |
str | Path |
Output folder for aligned mask tiles |
mask_key |
str | None |
NPZ mask key; auto-detected if None |
nodata |
int |
Fill for areas with no annotation (default 0) |
skip_empty |
bool |
Skip grid cells with no coverage |
overlapping |
str |
How to resolve overlapping (contradicting) annotation sets: "first" (default), "last", or "vote" (per-pixel majority) |
Returns the written tile names (tile_NNNNN.tif), pairing with download_grid imagery by index.
Overlaps & contradictions. Where several annotation sets cover the same pixel with different classes,
overlapping="vote"takes the per-pixel majority across all overlapping tiles;"first"/"last"keep the first/last tile's value.
realign_annotations_from_shp(shp_path, res, patch, annotations_dir, manifest, out_dir, mask_key=None, nodata=0, skip_empty=True, overlapping="first") -> list[str]
One call: build_grid(shp_path, res, patch) then realign_annotations_to_grid. res should
match the annotation resolution. overlapping resolves contradicting overlaps (see above).
Quick QA plots. The web map helpers live in tytonai_utils.webmap; the image/mask helper
lives in tytonai_utils.viz. matplotlib comes from the viz (or webmap) extra.
Downscaled overview mosaic of downloaded .tif tiles, placed at their real geo-extent.
Handles RGB (3+ bands) and single-band greyscale. Save with out_png. (See Feature 1.)
Draw the tile grid over the study-area outline — no download needed (cell size in the title is read from the grid). (See Feature 1.)
viz.plot_image_mask_pairs(annotations_dir, manifest, indexes=None, n=6, out_png=None, mask_dir=None, image_key=None, mask_key=None, bands=(0,1,2), rgb_keys=("RED","GREEN","BLUE"), show_dsm=False, dsm_cmap="terrain", class_names=None, cmap="tab20", max_rows=3, seed=0) -> Figure
Plot imagery tiles next to their annotation masks (the .npz pairs from Feature 2) for
visual QA, with a class legend. The imagery npz stores bands as separate keys, so RGB is
composed from rgb_keys (RED/GREEN/BLUE); set show_dsm=True to add a DSM panel per
sample. Select specific indexes, else n random tiles (seeded). Layout is capped at
max_rows rows and grows columns. Mask classes get consistent colours across panels.
| Param | Type | Description |
|---|---|---|
annotations_dir |
str | Path |
Folder of downloaded imagery + mask .npz |
manifest |
list[dict] | str | Path |
Tile list (read_manifest) or path to dataset.json |
indexes |
list[int] | None |
Tiles to show; None → n random |
n |
int |
Number of random tiles when indexes is None |
out_png |
str | Path | None |
Save the figure if given |
mask_dir |
str | Path | None |
Folder to load masks from; defaults to annotations_dir. Point at a rolled-up folder (e.g. annotations_rnd7) to view remapped masks against the original imagery |
image_key / mask_key |
str | None |
NPZ array keys; auto-detected if None |
bands |
tuple[int,...] |
Channel indices to render as RGB (when imagery is a single 3D array) |
rgb_keys |
tuple[str,str,str] |
NPZ keys to stack into RGB (default RED/GREEN/BLUE) |
show_dsm |
bool |
Add a DSM panel (from the DSM key) per sample |
dsm_cmap |
str |
Colormap for the DSM panel |
class_names |
dict[int,str] | None |
Label the legend by name; defaults to the org class names (CLASS_NAMES), which also cover rolled-up ids. Pass RND_NAMES_7CLASS etc. to override |
cmap |
str |
Colormap for the mask classes (categorical) |
max_rows |
int |
Max rows; samples wrap into more columns |
seed |
int |
RNG seed for reproducible random selection |
Returns the matplotlib Figure. Needs the manifest because each tile's image and mask
are separate npz files referenced by it. For grid-aligned .tif tiles paired by filename, use
plot_image_mask_tiles instead.
viz.plot_image_mask_tiles(image_dir, mask_dir, indexes=None, n=6, out_png=None, bands=(1,2,3), class_names=None, cmap="tab20", max_rows=3, seed=0) -> Figure
Plot imagery .tif next to mask .tif, paired by filename across two folders — no
manifest needed. Built for grid-aligned outputs: image_dir from download_grid, mask_dir
from realign_annotations_to_grid (both write tile_NNNNN.tif). Pairs are the tiles present in
both folders. Same layout + legend as plot_image_mask_pairs. Needs rasterio (webmap extra).
| Param | Type | Description |
|---|---|---|
image_dir |
str | Path |
Folder of RGB imagery .tif (e.g. download_grid output) |
mask_dir |
str | Path |
Folder of mask .tif (e.g. realign_annotations_to_grid output) |
indexes |
list[int] | None |
Positions in the sorted paired list; None → n random |
n |
int |
Number of random pairs when indexes is None |
out_png |
str | Path | None |
Save the figure if given |
bands |
tuple[int,...] |
1-based GeoTIFF bands for RGB |
class_names |
dict[int,str] | None |
Label the legend by name; else by id |
cmap |
str |
Colormap for the mask classes |
max_rows |
int |
Max rows; samples wrap into more columns |
seed |
int |
RNG seed for reproducible random selection |
Returns the matplotlib Figure.
Shared boto3 client for the S3 API (used by features 2–4; feature 1 uses GDAL directly).
Return a boto3 S3 client configured for the tytonai endpoint. Pass env explicitly to
stay testable; defaults to os.environ. Forces path-style addressing when
AWS_VIRTUAL_HOSTING is FALSE.
from dotenv import load_dotenv
from tytonai_utils.s3 import make_s3_client
load_dotenv()
s3 = make_s3_client()
s3.list_objects_v2(Bucket="...", MaxKeys=5)RuntimeError: main thread is not in main loop / Tcl_AsyncDelete … core dumped during
download_grid. matplotlib's interactive Tk backend is not thread-safe and clashes
with the threaded downloads. Handled automatically: importing tytonai_utils selects the
non-interactive Agg backend, so plots save to PNG and downloads never crash — nothing to
configure. Just import tytonai_utils (or any from tytonai_utils...) before importing
matplotlib.pyplot yourself. To override, set MPLBACKEND before launching Python.
git clone https://github.com/clement-tyton/tytonai_utils.git
cd tytonai_utils
uv sync --extra all --extra test
uv run pytestEach module has a runnable RUN block at the bottom (if __name__ == "__main__":)
designed to be stepped through line by line in VSCode.