Skip to content

rslearn.models.clip

clip

OpenAI CLIP models.

CLIP

Bases: FeatureExtractor

CLIP image encoder.

Source code in rslearn/models/clip.py 
10
11
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
class CLIP(FeatureExtractor):
    """CLIP image encoder."""

    def __init__(
        self,
        model_name: str,
    ):
        """Instantiate a new CLIP instance.

        Args:
            model_name: the model name like "openai/clip-vit-large-patch14-336".
        """
        super().__init__()

        self.processor = AutoProcessor.from_pretrained(model_name)  # nosec
        model = AutoModelForZeroShotImageClassification.from_pretrained(model_name)  # nosec
        self.encoder = model.vision_model

        # Get number of features and token map size from encoder attributes.
        self.num_features = self.encoder.post_layernorm.normalized_shape[0]
        crop_size = self.processor.image_processor.crop_size
        stride = self.encoder.embeddings.patch_embedding.stride
        self.height = crop_size["height"] // stride[0]
        self.width = crop_size["width"] // stride[1]

    def forward(self, context: ModelContext) -> FeatureMaps:
        """Compute outputs from the backbone.

        Args:
            context: the model context. Input dicts must include "image" key containing
                the image to process. The images should have values 0-255.

        Returns:
            a FeatureMaps with one feature map from the ViT, which is always Bx24x24x1024.
        """
        inputs = context.inputs
        device = inputs[0]["image"].image.device
        clip_inputs = self.processor(
            images=[
                inp["image"].single_ts_to_chw_tensor().cpu().numpy().transpose(1, 2, 0)
                for inp in inputs
            ],
            return_tensors="pt",
            padding=True,
        )
        pixel_values = clip_inputs["pixel_values"].to(device)
        output = self.encoder(pixel_values=pixel_values)
        # Ignore class token output which is before the patch tokens.
        image_features = output.last_hidden_state[:, 1:, :]
        batch_size = image_features.shape[0]

        # 576x1024 -> HxWxC
        return FeatureMaps(
            [
                image_features.reshape(
                    batch_size, self.height, self.width, self.num_features
                ).permute(0, 3, 1, 2)
            ]
        )

forward 

forward(context: ModelContext) -> FeatureMaps

Compute outputs from the backbone.

Parameters:

Name Type Description Default
context ModelContext

the model context. Input dicts must include "image" key containing the image to process. The images should have values 0-255.

 required

Returns:

Type Description
FeatureMaps

a FeatureMaps with one feature map from the ViT, which is always Bx24x24x1024.
Source code in rslearn/models/clip.py 
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
def forward(self, context: ModelContext) -> FeatureMaps:
    """Compute outputs from the backbone.

    Args:
        context: the model context. Input dicts must include "image" key containing
            the image to process. The images should have values 0-255.

    Returns:
        a FeatureMaps with one feature map from the ViT, which is always Bx24x24x1024.
    """
    inputs = context.inputs
    device = inputs[0]["image"].image.device
    clip_inputs = self.processor(
        images=[
            inp["image"].single_ts_to_chw_tensor().cpu().numpy().transpose(1, 2, 0)
            for inp in inputs
        ],
        return_tensors="pt",
        padding=True,
    )
    pixel_values = clip_inputs["pixel_values"].to(device)
    output = self.encoder(pixel_values=pixel_values)
    # Ignore class token output which is before the patch tokens.
    image_features = output.last_hidden_state[:, 1:, :]
    batch_size = image_features.shape[0]

    # 576x1024 -> HxWxC
    return FeatureMaps(
        [
            image_features.reshape(
                batch_size, self.height, self.width, self.num_features
            ).permute(0, 3, 1, 2)
        ]
    )