Foundation Models API

Auto-generated from source — see Foundation Models for narrative descriptions.

Prithvi

keras_climate.foundation.prithvi.PrithviEncoder

PrithviEncoder(img_size=224, patch_size=16, num_frames=3, tubelet_size=1, in_chans=6, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0, name='prithvi_encoder', **kwargs)

Bases: Model

Source code in keras_climate/foundation/prithvi.py

def __init__(
    self,
    img_size=224,
    patch_size=16,
    num_frames=3,
    tubelet_size=1,
    in_chans=6,
    embed_dim=768,
    depth=12,
    num_heads=12,
    mlp_ratio=4.0,
    name="prithvi_encoder",
    **kwargs,
):
    super().__init__(name=name, **kwargs)
    self.embed_dim = embed_dim
    self.num_frames = num_frames
    self.tubelet_size = tubelet_size
    self.grid_size = img_size // patch_size
    self.t_grid = num_frames // tubelet_size

    self.patch_embed = PatchEmbed3D(
        patch_size, tubelet_size, embed_dim, name="patch_embed"
    )
    self.cls_token = self.add_weight(
        shape=(1, 1, embed_dim),
        initializer="zeros",
        trainable=True,
        name="cls_token",
    )

    pos = get_3d_sincos_pos_embed(
        embed_dim, (self.t_grid, self.grid_size, self.grid_size), add_cls_token=True
    )
    self.pos_embed = self.add_weight(
        shape=pos.shape,
        initializer=keras.initializers.Constant(pos),
        trainable=False,
        name="pos_embed",
    )

    self.blocks = [
        TransformerEncoderBlock(embed_dim, num_heads, mlp_ratio, name=f"block{i}")
        for i in range(depth)
    ]
    self.norm = layers.LayerNormalization(epsilon=1e-6, name="norm")

keras_climate.foundation.prithvi.PrithviClassifier

PrithviClassifier(variant='prithvi_100m', img_size=224, patch_size=16, num_frames=3, in_chans=6, num_classes=10, name='prithvi_classifier')

Source code in keras_climate/foundation/prithvi.py

def PrithviClassifier(
    variant="prithvi_100m",
    img_size=224,
    patch_size=16,
    num_frames=3,
    in_chans=6,
    num_classes=10,
    name="prithvi_classifier",
):
    cfg = PRITHVI_CONFIGS[variant]
    inputs = keras.Input(
        shape=(num_frames, img_size, img_size, in_chans), name="hls_stack"
    )
    encoder = PrithviEncoder(
        img_size,
        patch_size,
        num_frames,
        1,
        in_chans,
        cfg["embed_dim"],
        cfg["depth"],
        cfg["num_heads"],
        name="encoder",
    )
    cls_token = encoder(inputs, return_all_tokens=False)
    outputs = layers.Dense(num_classes, name="head")(cls_token)
    return keras.Model(inputs, outputs, name=name)

keras_climate.foundation.prithvi.PrithviSegmenter

PrithviSegmenter(variant='prithvi_100m', img_size=224, patch_size=16, num_frames=3, in_chans=6, num_classes=2, name='prithvi_segmenter')

Source code in keras_climate/foundation/prithvi.py

def PrithviSegmenter(
    variant="prithvi_100m",
    img_size=224,
    patch_size=16,
    num_frames=3,
    in_chans=6,
    num_classes=2,
    name="prithvi_segmenter",
):
    cfg = PRITHVI_CONFIGS[variant]
    inputs = keras.Input(
        shape=(num_frames, img_size, img_size, in_chans), name="hls_stack"
    )
    encoder = PrithviEncoder(
        img_size,
        patch_size,
        num_frames,
        1,
        in_chans,
        cfg["embed_dim"],
        cfg["depth"],
        cfg["num_heads"],
        name="encoder",
    )
    tokens = encoder(inputs, return_all_tokens=True)
    patch_tokens = layers.Lambda(lambda t: t[:, 1:, :], name="drop_cls")(tokens)

    grid = img_size // patch_size
    t_grid = num_frames
    x = layers.Reshape((t_grid, grid, grid, cfg["embed_dim"]), name="to_grid")(
        patch_tokens
    )
    x = layers.Lambda(lambda t: ops.mean(t, axis=1), name="temporal_pool")(x)

    for i, f in enumerate([256, 128, 64]):
        x = layers.Conv2DTranspose(f, 3, strides=2, padding="same", name=f"up{i}")(x)
        x = layers.BatchNormalization(name=f"up_bn{i}")(x)
        x = layers.Activation("relu", name=f"up_relu{i}")(x)

    x = layers.Lambda(
        lambda t: ops.image.resize(t, (img_size, img_size), interpolation="bilinear"),
        name="resize_to_input",
    )(x)
    outputs = layers.Conv2D(num_classes, 1, name="seg_head")(x)
    return keras.Model(inputs, outputs, name=name)

keras_climate.foundation.prithvi.PRITHVI_CONFIGS module-attribute

PRITHVI_CONFIGS = {'prithvi_100m': dict(embed_dim=768, depth=12, num_heads=12, decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16), 'prithvi_300m': dict(embed_dim=1024, depth=24, num_heads=16, decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16), 'prithvi_eo_v2_300m': dict(embed_dim=1024, depth=24, num_heads=16, decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16, patch_size=16), 'prithvi_eo_v2_600m': dict(embed_dim=1280, depth=32, num_heads=16, decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16, patch_size=14)}

Clay

keras_climate.foundation.clay.ClayEncoder

ClayEncoder(img_size=224, patch_size=8, embed_dim=768, depth=12, num_heads=12, dim_head=64, mlp_ratio=4.0, wave_dim=128, num_latent_tokens=128, name='clay_encoder', **kwargs)

Bases: Model

Source code in keras_climate/foundation/clay.py

def __init__(
    self,
    img_size=224,
    patch_size=8,
    embed_dim=768,
    depth=12,
    num_heads=12,
    dim_head=64,
    mlp_ratio=4.0,
    wave_dim=128,
    num_latent_tokens=128,
    name="clay_encoder",
    **kwargs,
):
    super().__init__(name=name, **kwargs)
    self.embed_dim = embed_dim
    self.patch_size = patch_size
    self.grid = img_size // patch_size

    self.cls_token = self.add_weight(
        shape=(1, 1, embed_dim), initializer="random_normal", name="cls_token"
    )
    self.patch_embedding = DynamicEmbedding(
        wave_dim, num_latent_tokens, patch_size, embed_dim, name="patch_embedding"
    )
    self.blocks = [
        ClayTransformerBlock(
            embed_dim,
            num_heads,
            dim_head,
            int(embed_dim * mlp_ratio),
            name=f"block{i}",
        )
        for i in range(depth)
    ]
    self.norm = layers.LayerNormalization(epsilon=_LN_EPS, name="norm")

keras_climate.foundation.clay.ClayClassifier

ClayClassifier(img_size=224, patch_size=8, embed_dim=768, depth=12, num_heads=12, num_bands=10, num_classes=10, name='clay_classifier')

Source code in keras_climate/foundation/clay.py

def ClayClassifier(
    img_size=224,
    patch_size=8,
    embed_dim=768,
    depth=12,
    num_heads=12,
    num_bands=10,
    num_classes=10,
    name="clay_classifier",
):
    pixels_in = keras.Input((img_size, img_size, num_bands), name="pixels")
    time_latlon_in = keras.Input((8,), name="time_latlon")
    waves_in = keras.Input((num_bands,), batch_size=1, name="waves")

    encoder = ClayEncoder(
        img_size, patch_size, embed_dim, depth, num_heads, name="encoder"
    )
    tokens = encoder(
        {
            "pixels": pixels_in,
            "waves": layers.Lambda(lambda t: t[0])(waves_in),
            "time_latlon": time_latlon_in,
        }
    )
    cls_token = layers.Lambda(lambda t: t[:, 0], name="take_cls")(tokens)
    outputs = layers.Dense(num_classes, name="head")(cls_token)
    return keras.Model([pixels_in, time_latlon_in, waves_in], outputs, name=name)

CROMA

keras_climate.foundation.croma.CROMA

CROMA(img_size=120, patch_size=8, size='base', sar_chans=2, optical_chans=12, name='croma')

Source code in keras_climate/foundation/croma.py

def CROMA(
    img_size=120,
    patch_size=8,
    size="base",
    sar_chans=2,
    optical_chans=12,
    name="croma",
):
    dim, depth = {"base": (768, 12), "large": (1024, 24)}[size]
    num_heads = 16
    grid_size = img_size // patch_size
    attn_bias_np = get_2d_alibi(num_heads, grid_size)

    sar_in = keras.Input((img_size, img_size, sar_chans), name="sar")
    opt_in = keras.Input((img_size, img_size, optical_chans), name="optical")

    attn_bias = ops.convert_to_tensor(attn_bias_np)

    sar_encoder = ModalityEncoder(
        dim, depth // 2, sar_chans, patch_size, num_heads, name="s1_encoder"
    )
    opt_encoder = ModalityEncoder(
        dim, depth, optical_chans, patch_size, num_heads, name="s2_encoder"
    )

    sar_tokens = sar_encoder(sar_in, attn_bias)
    opt_tokens = opt_encoder(opt_in, attn_bias)

    fusion = CrossAttentionFusion(dim, num_heads, depth // 2, name="cross_encoder")
    joint_tokens = fusion(sar_tokens, opt_tokens, attn_bias)

    sar_pool = layers.GlobalAveragePooling1D(name="sar_pool")(sar_tokens)
    opt_pool = layers.GlobalAveragePooling1D(name="optical_pool")(opt_tokens)
    sar_repr = gap_ffn(dim, name="GAP_FFN_s1")(sar_pool)
    opt_repr = gap_ffn(dim, name="GAP_FFN_s2")(opt_pool)

    return keras.Model(
        [sar_in, opt_in],
        {"sar_repr": sar_repr, "optical_repr": opt_repr, "joint_tokens": joint_tokens},
        name=name,
    )

keras_climate.foundation.croma.ModalityEncoder

ModalityEncoder(dim=768, depth=12, in_chans=2, patch_size=8, num_heads=16, mlp_mult=4, name='modality_encoder', **kwargs)

Bases: Model

Source code in keras_climate/foundation/croma.py

def __init__(
    self,
    dim=768,
    depth=12,
    in_chans=2,
    patch_size=8,
    num_heads=16,
    mlp_mult=4,
    name="modality_encoder",
    **kwargs,
):
    super().__init__(name=name, **kwargs)
    self.dim = dim
    self.patch_size = patch_size
    self.linear_input = layers.Dense(dim, name="linear_input")
    self.blocks = [
        CromaSelfBlock(dim, num_heads, mlp_mult, name=f"block{i}")
        for i in range(depth)
    ]
    self.norm_out = layers.LayerNormalization(epsilon=_LN_EPS, name="norm_out")

keras_climate.foundation.croma.CrossAttentionFusion

CrossAttentionFusion(dim, num_heads=16, depth=6, mlp_mult=4, name='cross_encoder', **kwargs)

Bases: Model

Source code in keras_climate/foundation/croma.py

def __init__(
    self, dim, num_heads=16, depth=6, mlp_mult=4, name="cross_encoder", **kwargs
):
    super().__init__(name=name, **kwargs)
    self.blocks = [
        CromaCrossBlock(dim, num_heads, mlp_mult, name=f"block{i}")
        for i in range(depth)
    ]
    self.norm_out = layers.LayerNormalization(epsilon=_LN_EPS, name="norm_out")

AnySat (this repo's design)

keras_climate.foundation.anysat.AnySatEncoder

AnySatEncoder(modality_specs, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0, name='anysat_encoder', **kwargs)

Bases: Model

Source code in keras_climate/foundation/anysat.py

def __init__(
    self,
    modality_specs,
    embed_dim=768,
    depth=12,
    num_heads=12,
    mlp_ratio=4.0,
    name="anysat_encoder",
    **kwargs,
):
    super().__init__(name=name, **kwargs)
    self.modality_specs = modality_specs
    self.embed_dim = embed_dim

    self.patch_embeds = {
        m: ModalityPatchEmbed(
            embed_dim, spec["patch_size_px"], spec["gsd_m"], name=f"embed_{m}"
        )
        for m, spec in modality_specs.items()
    }
    self.modality_tokens = {
        m: self.add_weight(
            shape=(1, 1, embed_dim), initializer="zeros", name=f"modtok_{m}"
        )
        for m in modality_specs
    }
    self.pos_encoding = ContinuousPositionEncoding(embed_dim, name="pos_encoding")
    self.cls_token = self.add_weight(
        shape=(1, 1, embed_dim), initializer="zeros", name="cls_token"
    )

    self.blocks = [
        TransformerEncoderBlock(embed_dim, num_heads, mlp_ratio, name=f"block{i}")
        for i in range(depth)
    ]
    self.norm = layers.LayerNormalization(epsilon=1e-6, name="norm")

keras_climate.foundation.anysat.AnySatClassifier

AnySatClassifier(modality_specs, embed_dim=768, depth=12, num_heads=12, input_shapes=None, num_classes=10, name='anysat_classifier')

Source code in keras_climate/foundation/anysat.py

def AnySatClassifier(
    modality_specs,
    embed_dim=768,
    depth=12,
    num_heads=12,
    input_shapes=None,
    num_classes=10,
    name="anysat_classifier",
):
    inputs = {m: keras.Input(shape=shape, name=m) for m, shape in input_shapes.items()}
    encoder = AnySatEncoder(modality_specs, embed_dim, depth, num_heads, name="encoder")
    tokens = encoder(inputs)
    cls_token = layers.Lambda(lambda t: t[:, 0], name="take_cls")(tokens)
    outputs = layers.Dense(num_classes, name="head")(cls_token)
    return keras.Model(inputs, outputs, name=name)

keras_climate.foundation.anysat.ModalityPatchEmbed

ModalityPatchEmbed(embed_dim, patch_size_px, gsd_m, **kwargs)

Bases: Layer

Source code in keras_climate/foundation/anysat.py

def __init__(self, embed_dim, patch_size_px, gsd_m, **kwargs):
    super().__init__(**kwargs)
    self.embed_dim = embed_dim
    self.patch_size_px = patch_size_px
    self.gsd_m = gsd_m

AnySat (official release port)

keras_climate.foundation.anysat_release.AnySatRelease

AnySatRelease(modalities, input_shapes, scale, size='base', embed_dim=None, depth=None, num_heads=None, mlp_ratio=4.0, output='patch', name='anysat_release', **kwargs)

Bases: Model

Source code in keras_climate/foundation/anysat_release.py

def __init__(
    self,
    modalities,
    input_shapes,
    scale,
    size="base",
    embed_dim=None,
    depth=None,
    num_heads=None,
    mlp_ratio=4.0,
    output="patch",
    name="anysat_release",
    **kwargs,
):
    super().__init__(name=name, **kwargs)
    cfg = ANYSAT_CONFIGS[size]
    self.embed_dim = embed_dim or cfg["embed_dim"]
    self.depth = depth or cfg["depth"]
    self.num_heads = num_heads or cfg["num_heads"]
    self.mlp_ratio = mlp_ratio
    self.modalities = list(modalities)
    self.input_shapes = dict(input_shapes)
    self.scale = scale
    self.output_mode = output
    if output not in ("tile", "patch"):
        raise ValueError(
            "output must be 'tile' or 'patch' (dense/subpatch output is not implemented)"
        )

    proj_cfgs = anysat_projector_configs(self.embed_dim)
    self.projectors = {}
    for m in self.modalities:
        pc = proj_cfgs[m]
        if pc["kind"] == "image":
            self.projectors[m] = AnySatImageProjector(
                self.embed_dim,
                pc["patch_size"],
                pc["in_chans"],
                pc["resolution"],
                scale,
                bias=False,
                name=f"projector_{m}",
            )
        else:
            self.projectors[m] = AnySatTimeSeriesProjector(
                self.embed_dim,
                pc["in_channels"],
                pc["T"],
                pc["in_norm"],
                pc["reduce_scale"],
                scale,
                n_head=_TS_LTAE_HEAD["n_head"],
                d_k=_TS_LTAE_HEAD["d_k"],
                mlp_in=pc["mlp_in"],
                name=f"projector_{m}",
            )

    self.local_encoder = AnySatLocalEncoder(
        self.embed_dim,
        self.depth,
        self.num_heads,
        mlp_ratio,
        name="spatial_encoder",
    )
    self.cls_token = self.add_weight(
        shape=(1, 1, self.embed_dim), initializer="zeros", name="cls_token"
    )
    self.blocks = [
        AnySatPlainBlock(
            self.embed_dim, self.num_heads, mlp_ratio, name=f"block{i}"
        )
        for i in range(self.depth)
    ]
    self.cross_pool = AnySatCrossPoolingBlock(
        self.embed_dim, self.num_heads, mlp_ratio, name="block_cross"
    )

    self._build_meta()

keras_climate.foundation.anysat_release.ANYSAT_CONFIGS module-attribute

ANYSAT_CONFIGS = {'tiny': dict(embed_dim=256, depth=2, num_heads=4), 'small': dict(embed_dim=512, depth=4, num_heads=8), 'base': dict(embed_dim=768, depth=6, num_heads=12)}

keras_climate.foundation.anysat_release.ANYSAT_MODALITIES module-attribute

ANYSAT_MODALITIES = ['aerial', 'aerial-flair', 'spot', 'naip', 's2', 's1-asc', 's1', 'alos', 'l7', 'l8', 'modis']