Skip to content

Add lat/lon support to OlmoEarth finetuning wrapper#574

Open
pjreddie wants to merge 1 commit intomasterfrom
feature/latlon-encoding
Open

Add lat/lon support to OlmoEarth finetuning wrapper#574
pjreddie wants to merge 1 commit intomasterfrom
feature/latlon-encoding

Conversation

@pjreddie
Copy link
Copy Markdown
Contributor

Summary

  • Adds use_latlon: bool parameter to OlmoEarth wrapper
  • When enabled, computes lat/lon from SampleMetadata crop center + CRS projection via rasterio.warp.transform
  • Sets latlon on MaskedOlmoEarthSample so the encoder applies geographic encoding during finetuning
  • Requires pretrained model trained with use_latlon_encoding=True

Related PRs

Test plan

  • Unit tests: WGS84 identity transform, UTM non-trivial transform, batch processing
  • Integration tests: full forward pass with use_latlon=True, backward compat with use_latlon=False
  • All 17 existing OlmoEarth tests still pass

🤖 Generated with Claude Code

Adds `use_latlon` parameter to OlmoEarth. When enabled, computes
lat/lon from SampleMetadata crop center + CRS projection and passes
it to the encoder for geographic encoding during finetuning.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
from rasterio.crs import CRS
from rasterio.warp import transform

wgs84 = CRS.from_epsg(4326)
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

More consistent with rest of code to use rslearn tooling for this:

from rslearn.const import WGS84_PROJECTION
geom = STGeometry(meta.projection.crs, shapely.Point(cx, cy), None)
wgs84_geom = geom.to_projection(WGS84_PROJECTION)
latlons.append([wgs84_geom.shp.y, wgs84_geom.shp.x])

Returns:
Tensor of shape (B, 2) with (latitude, longitude) in degrees.
"""
from rasterio.crs import CRS
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Move imports to top of file.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants