Skip to content

rslearn.models.sam2_enc

sam2_enc

SegmentAnything2 encoders.

SAM2Encoder

Bases: FeatureExtractor

SAM2's image encoder.

Source code in rslearn/models/sam2_enc.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
class SAM2Encoder(FeatureExtractor):
    """SAM2's image encoder."""

    def __init__(self, model_identifier: str) -> None:
        """Initializes the SAM2Encoder with a specific model configuration and checkpoint.

        Args:
            model_identifier: Identifier for model type.
        """
        super().__init__()

        if "tiny" in model_identifier:
            model_cfg = "sam2_hiera_t.yaml"
            checkpoint_path = UPath(
                "gcs://rslearn-eai/artifacts/sam2/sam2_hiera_tiny.pt"
            )
            self.backbone_channels = [
                [4, 96],
                [8, 192],
                [16, 384],
                [32, 768],
            ]
        elif "small" in model_identifier:
            model_cfg = "sam2_hiera_s.yaml"
            checkpoint_path = UPath(
                "gcs://rslearn-eai/artifacts/sam2/sam2_hiera_small.pt"
            )
            self.backbone_channels = [
                [4, 96],
                [8, 192],
                [16, 384],
                [32, 768],
            ]
        elif "base" in model_identifier:
            model_cfg = "sam2_hiera_b+.yaml"
            checkpoint_path = UPath(
                "gcs://rslearn-eai/artifacts/sam2/sam2_hiera_base_plus.pt"
            )
            self.backbone_channels = [
                [4, 112],
                [8, 224],
                [16, 448],
                [32, 896],
            ]
        elif "large" in model_identifier:
            model_cfg = "sam2_hiera_l.yaml"
            checkpoint_path = UPath(
                "gcs://rslearn-eai/artifacts/sam2/sam2_hiera_large.pt"
            )
            self.backbone_channels = [
                [4, 144],
                [8, 288],
                [16, 576],
                [32, 1152],
            ]
        else:
            raise ValueError(f"Invalid model identifier: {model_identifier}")

        # Build the model and remove unnecessary components
        with checkpoint_path.open("rb") as f:
            self.model = build_sam2(model_cfg, f)
        self._remove_unused_modules()

        self.encoder = self.model.image_encoder.trunk

    def _remove_unused_modules(self) -> None:
        """Removes unused modules from the SAM2 model."""
        del self.model.sam_mask_decoder
        del self.model.sam_prompt_encoder
        del self.model.memory_encoder
        del self.model.memory_attention
        del self.model.mask_downsample
        del self.model.obj_ptr_tpos_proj
        del self.model.obj_ptr_proj
        del self.model.image_encoder.neck

    def forward(self, context: ModelContext) -> FeatureMaps:
        """Extract multi-scale features from a batch of images.

        Args:
            context: the model context. Input dicts must have a key 'image' containing
                the input for the SAM2 image encoder.

        Returns:
            feature maps from the encoder.
        """
        images = torch.stack(
            [inp["image"].single_ts_to_chw_tensor() for inp in context.inputs], dim=0
        )
        features = self.encoder(images)
        return FeatureMaps(features)

    def get_backbone_channels(self) -> list[list[int]]:
        """Returns the output channels of the encoder at different scales.

        Returns:
            List[List[int]]: List of downsample factors and corresponding channel counts at each scale.
        """
        return self.backbone_channels

forward 

forward(context: ModelContext) -> FeatureMaps

Extract multi-scale features from a batch of images.

Parameters:

Name Type Description Default
context ModelContext

the model context. Input dicts must have a key 'image' containing the input for the SAM2 image encoder.

 required

Returns:

Type Description
FeatureMaps

feature maps from the encoder.
Source code in rslearn/models/sam2_enc.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def forward(self, context: ModelContext) -> FeatureMaps:
    """Extract multi-scale features from a batch of images.

    Args:
        context: the model context. Input dicts must have a key 'image' containing
            the input for the SAM2 image encoder.

    Returns:
        feature maps from the encoder.
    """
    images = torch.stack(
        [inp["image"].single_ts_to_chw_tensor() for inp in context.inputs], dim=0
    )
    features = self.encoder(images)
    return FeatureMaps(features)

get_backbone_channels 

get_backbone_channels() -> list[list[int]]

Returns the output channels of the encoder at different scales.

Returns:

Type Description
list[list[int]]

List[List[int]]: List of downsample factors and corresponding channel counts at each scale.
Source code in rslearn/models/sam2_enc.py 
104
105
106
107
108
109
110
def get_backbone_channels(self) -> list[list[int]]:
    """Returns the output channels of the encoder at different scales.

    Returns:
        List[List[int]]: List of downsample factors and corresponding channel counts at each scale.
    """
    return self.backbone_channels