Upload 4 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+baklava.jpg filter=lfs diff=lfs merge=lfs -text
+cat.jpg filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,159 @@
+"""MobileCLIP2 Zero-Shot Classification Demo"""
+import torch
+import open_clip
+import gradio as gr
+from mobileclip.modules.common.mobileone import reparameterize_model
+
+################################################################################
+# Model Configuration
+################################################################################
+AVAILABLE_MODELS = {
+    "MobileCLIP2-B": ("MobileCLIP2-B", "dfndr2b"),
+    "MobileCLIP2-S0": ("MobileCLIP2-S0", "dfndr2b"),
+    "MobileCLIP2-S2": ("MobileCLIP2-S2", "dfndr2b"),
+    "MobileCLIP2-S3": ("MobileCLIP2-S3", "dfndr2b"),
+    "MobileCLIP2-S4": ("MobileCLIP2-S4", "dfndr2b"),
+    "MobileCLIP2-L-14": ("MobileCLIP2-L-14", "dfndr2b"),
+}
+
+# Cache for loaded models
+model_cache = {}
+
+################################################################################
+# Model Loading
+################################################################################
+def load_model(model_name):
+    """Load and cache MobileCLIP2 model"""
+    if model_name in model_cache:
+        return model_cache[model_name]
+    
+    model_id, pretrained = AVAILABLE_MODELS[model_name]
+    
+    # Create model and preprocessing transforms
+    model, _, preprocess = open_clip.create_model_and_transforms(
+        model_id, 
+        pretrained=pretrained
+    )
+    tokenizer = open_clip.get_tokenizer(model_id)
+    
+    # Reparameterize model for inference
+    model = reparameterize_model(model.eval())
+    
+    # Cache the model components
+    model_cache[model_name] = {
+        "model": model,
+        "preprocess": preprocess,
+        "tokenizer": tokenizer
+    }
+    
+    return model_cache[model_name]
+
+################################################################################
+# Inference
+################################################################################
+def classify_image(image, candidate_labels, model_name):
+    """
+    Classify image using selected MobileCLIP2 model
+    
+    Args:
+        image: PIL Image
+        candidate_labels: comma-separated string of labels
+        model_name: selected model from dropdown
+    
+    Returns:
+        Dictionary of label probabilities
+    """
+    if image is None:
+        return {}
+    
+    # Parse labels
+    labels = [label.strip() for label in candidate_labels.split(",") if label.strip()]
+    
+    if not labels:
+        return {}
+    
+    # Load model components
+    model_components = load_model(model_name)
+    model = model_components["model"]
+    preprocess = model_components["preprocess"]
+    tokenizer = model_components["tokenizer"]
+    
+    # Preprocess image
+    image_tensor = preprocess(image.convert('RGB')).unsqueeze(0)
+    
+    # Tokenize text
+    text_tokens = tokenizer(labels)
+    
+    # Run inference
+    # with torch.no_grad(), torch.cuda.amp.autocast():
+    with torch.no_grad():
+        image_features = model.encode_image(image_tensor)
+        text_features = model.encode_text(text_tokens)
+        
+        # Normalize features
+        image_features /= image_features.norm(dim=-1, keepdim=True)
+        text_features /= text_features.norm(dim=-1, keepdim=True)
+        
+        # Compute similarity and probabilities
+        text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)
+    
+    # Format output as dictionary
+    output = {labels[i]: float(text_probs[0][i]) for i in range(len(labels))}
+    
+    return output
+
+################################################################################
+# Gradio Interface
+################################################################################
+with gr.Blocks() as demo:
+    gr.Markdown("# MobileCLIP2 Zero-Shot Image Classification")
+    gr.Markdown(
+        "Classify images using MobileCLIP2 models. Select a model, upload an image, "
+        "and provide comma-separated class labels to get predictions."
+    )
+    
+    with gr.Row():
+        with gr.Column():
+            model_dropdown = gr.Dropdown(
+                choices=list(AVAILABLE_MODELS.keys()),
+                value="MobileCLIP2-S2",
+                label="Select MobileCLIP2 Model",
+                info="Choose which model to use for classification"
+            )
+            image_input = gr.Image(type="pil", label="Upload Image")
+            text_input = gr.Textbox(
+                label="Class Labels (comma separated)",
+                placeholder="e.g., a cat, a dog, a bird"
+            )
+            run_button = gr.Button("Classify", variant="primary")
+        
+        with gr.Column():
+            output_label = gr.Label(
+                label="Classification Results",
+                num_top_classes=5
+            )
+    
+    # Examples
+    examples = [
+        ["MobileCLIP2-S2", "./baklava.jpg", "dessert on a plate, a serving of baklava, a plate and spoon"],
+        ["MobileCLIP2-S2", "./cat.jpg", "a cat, two cats, three cats"],
+        ["MobileCLIP2-S2", "./cat.jpg", "two sleeping cats, two cats playing, three cats laying down"],
+    ]
+    
+    gr.Examples(
+        examples=examples,
+        inputs=[model_dropdown, image_input, text_input],
+        outputs=[output_label],
+        fn=classify_image,
+        cache_examples=False
+    )
+    
+    # Connect button
+    run_button.click(
+        fn=classify_image,
+        inputs=[image_input, text_input, model_dropdown],
+        outputs=[output_label]
+    )
+
+if __name__ == "__main__":
+    demo.launch()

mobileclip/modules/common/mobileone.py

@@ -0,0 +1,341 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (C) 2024 Apple Inc. All Rights Reserved.
+#
+from typing import Union, Tuple
+
+import copy
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+__all__ = ["MobileOneBlock", "reparameterize_model"]
+
+
+class SEBlock(nn.Module):
+    """Squeeze and Excite module.
+
+    Pytorch implementation of `Squeeze-and-Excitation Networks` -
+    https://arxiv.org/pdf/1709.01507.pdf
+    """
+
+    def __init__(self, in_channels: int, rd_ratio: float = 0.0625) -> None:
+        """Construct a Squeeze and Excite Module.
+
+        Args:
+            in_channels: Number of input channels.
+            rd_ratio: Input channel reduction ratio.
+        """
+        super(SEBlock, self).__init__()
+        self.reduce = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=int(in_channels * rd_ratio),
+            kernel_size=1,
+            stride=1,
+            bias=True,
+        )
+        self.expand = nn.Conv2d(
+            in_channels=int(in_channels * rd_ratio),
+            out_channels=in_channels,
+            kernel_size=1,
+            stride=1,
+            bias=True,
+        )
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        """Apply forward pass."""
+        b, c, h, w = inputs.size()
+        x = F.avg_pool2d(inputs, kernel_size=[h, w])
+        x = self.reduce(x)
+        x = F.relu(x)
+        x = self.expand(x)
+        x = torch.sigmoid(x)
+        x = x.view(-1, c, 1, 1)
+        return inputs * x
+
+
+class MobileOneBlock(nn.Module):
+    """MobileOne building block.
+
+    This block has a multi-branched architecture at train-time
+    and plain-CNN style architecture at inference time
+    For more details, please refer to our paper:
+    `An Improved One millisecond Mobile Backbone` -
+    https://arxiv.org/pdf/2206.04040.pdf
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        padding: int = 0,
+        dilation: int = 1,
+        groups: int = 1,
+        inference_mode: bool = False,
+        use_se: bool = False,
+        use_act: bool = True,
+        use_scale_branch: bool = True,
+        num_conv_branches: int = 1,
+        activation: nn.Module = nn.GELU(),
+    ) -> None:
+        """Construct a MobileOneBlock module.
+
+        Args:
+            in_channels: Number of channels in the input.
+            out_channels: Number of channels produced by the block.
+            kernel_size: Size of the convolution kernel.
+            stride: Stride size.
+            padding: Zero-padding size.
+            dilation: Kernel dilation factor.
+            groups: Group number.
+            inference_mode: If True, instantiates model in inference mode.
+            use_se: Whether to use SE-ReLU activations.
+            use_act: Whether to use activation. Default: ``True``
+            use_scale_branch: Whether to use scale branch. Default: ``True``
+            num_conv_branches: Number of linear conv branches.
+        """
+        super(MobileOneBlock, self).__init__()
+        self.inference_mode = inference_mode
+        self.groups = groups
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.kernel_size = kernel_size
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_conv_branches = num_conv_branches
+
+        # Check if SE-ReLU is requested
+        if use_se:
+            self.se = SEBlock(out_channels)
+        else:
+            self.se = nn.Identity()
+
+        if use_act:
+            self.activation = activation
+        else:
+            self.activation = nn.Identity()
+
+        if inference_mode:
+            self.reparam_conv = nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding,
+                dilation=dilation,
+                groups=groups,
+                bias=True,
+            )
+        else:
+            # Re-parameterizable skip connection
+            self.rbr_skip = (
+                nn.BatchNorm2d(num_features=in_channels)
+                if out_channels == in_channels and stride == 1
+                else None
+            )
+
+            # Re-parameterizable conv branches
+            if num_conv_branches > 0:
+                rbr_conv = list()
+                for _ in range(self.num_conv_branches):
+                    rbr_conv.append(
+                        self._conv_bn(kernel_size=kernel_size, padding=padding)
+                    )
+                self.rbr_conv = nn.ModuleList(rbr_conv)
+            else:
+                self.rbr_conv = None
+
+            # Re-parameterizable scale branch
+            self.rbr_scale = None
+            if not isinstance(kernel_size, int):
+                kernel_size = kernel_size[0]
+            if (kernel_size > 1) and use_scale_branch:
+                self.rbr_scale = self._conv_bn(kernel_size=1, padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply forward pass."""
+        # Inference mode forward pass.
+        if self.inference_mode:
+            return self.activation(self.se(self.reparam_conv(x)))
+
+        # Multi-branched train-time forward pass.
+        # Skip branch output
+        identity_out = 0
+        if self.rbr_skip is not None:
+            identity_out = self.rbr_skip(x)
+
+        # Scale branch output
+        scale_out = 0
+        if self.rbr_scale is not None:
+            scale_out = self.rbr_scale(x)
+
+        # Other branches
+        out = scale_out + identity_out
+        if self.rbr_conv is not None:
+            for ix in range(self.num_conv_branches):
+                out += self.rbr_conv[ix](x)
+
+        return self.activation(self.se(out))
+
+    def reparameterize(self):
+        """Following works like `RepVGG: Making VGG-style ConvNets Great Again` -
+        https://arxiv.org/pdf/2101.03697.pdf. We re-parameterize multi-branched
+        architecture used at training time to obtain a plain CNN-like structure
+        for inference.
+        """
+        if self.inference_mode:
+            return
+        kernel, bias = self._get_kernel_bias()
+        self.reparam_conv = nn.Conv2d(
+            in_channels=self.in_channels,
+            out_channels=self.out_channels,
+            kernel_size=self.kernel_size,
+            stride=self.stride,
+            padding=self.padding,
+            dilation=self.dilation,
+            groups=self.groups,
+            bias=True,
+        )
+        self.reparam_conv.weight.data = kernel
+        self.reparam_conv.bias.data = bias
+
+        # Delete un-used branches
+        for para in self.parameters():
+            para.detach_()
+        self.__delattr__("rbr_conv")
+        self.__delattr__("rbr_scale")
+        if hasattr(self, "rbr_skip"):
+            self.__delattr__("rbr_skip")
+
+        self.inference_mode = True
+
+    def _get_kernel_bias(self) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Method to obtain re-parameterized kernel and bias.
+        Reference: https://github.com/DingXiaoH/RepVGG/blob/main/repvgg.py#L83
+
+        Returns:
+            Tuple of (kernel, bias) after fusing branches.
+        """
+        # get weights and bias of scale branch
+        kernel_scale = 0
+        bias_scale = 0
+        if self.rbr_scale is not None:
+            kernel_scale, bias_scale = self._fuse_bn_tensor(self.rbr_scale)
+            # Pad scale branch kernel to match conv branch kernel size.
+            pad = self.kernel_size // 2
+            kernel_scale = torch.nn.functional.pad(kernel_scale, [pad, pad, pad, pad])
+
+        # get weights and bias of skip branch
+        kernel_identity = 0
+        bias_identity = 0
+        if self.rbr_skip is not None:
+            kernel_identity, bias_identity = self._fuse_bn_tensor(self.rbr_skip)
+
+        # get weights and bias of conv branches
+        kernel_conv = 0
+        bias_conv = 0
+        if self.rbr_conv is not None:
+            for ix in range(self.num_conv_branches):
+                _kernel, _bias = self._fuse_bn_tensor(self.rbr_conv[ix])
+                kernel_conv += _kernel
+                bias_conv += _bias
+
+        kernel_final = kernel_conv + kernel_scale + kernel_identity
+        bias_final = bias_conv + bias_scale + bias_identity
+        return kernel_final, bias_final
+
+    def _fuse_bn_tensor(
+        self, branch: Union[nn.Sequential, nn.BatchNorm2d]
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Method to fuse batchnorm layer with preceeding conv layer.
+        Reference: https://github.com/DingXiaoH/RepVGG/blob/main/repvgg.py#L95
+
+        Args:
+            branch: Sequence of ops to be fused.
+
+        Returns:
+            Tuple of (kernel, bias) after fusing batchnorm.
+        """
+        if isinstance(branch, nn.Sequential):
+            kernel = branch.conv.weight
+            running_mean = branch.bn.running_mean
+            running_var = branch.bn.running_var
+            gamma = branch.bn.weight
+            beta = branch.bn.bias
+            eps = branch.bn.eps
+        else:
+            assert isinstance(branch, nn.BatchNorm2d)
+            if not hasattr(self, "id_tensor"):
+                input_dim = self.in_channels // self.groups
+
+                kernel_size = self.kernel_size
+                if isinstance(self.kernel_size, int):
+                    kernel_size = (self.kernel_size, self.kernel_size)
+
+                kernel_value = torch.zeros(
+                    (self.in_channels, input_dim, kernel_size[0], kernel_size[1]),
+                    dtype=branch.weight.dtype,
+                    device=branch.weight.device,
+                )
+                for i in range(self.in_channels):
+                    kernel_value[
+                        i, i % input_dim, kernel_size[0] // 2, kernel_size[1] // 2
+                    ] = 1
+                self.id_tensor = kernel_value
+            kernel = self.id_tensor
+            running_mean = branch.running_mean
+            running_var = branch.running_var
+            gamma = branch.weight
+            beta = branch.bias
+            eps = branch.eps
+        std = (running_var + eps).sqrt()
+        t = (gamma / std).reshape(-1, 1, 1, 1)
+        return kernel * t, beta - running_mean * gamma / std
+
+    def _conv_bn(self, kernel_size: int, padding: int) -> nn.Sequential:
+        """Helper method to construct conv-batchnorm layers.
+
+        Args:
+            kernel_size: Size of the convolution kernel.
+            padding: Zero-padding size.
+
+        Returns:
+            Conv-BN module.
+        """
+        mod_list = nn.Sequential()
+        mod_list.add_module(
+            "conv",
+            nn.Conv2d(
+                in_channels=self.in_channels,
+                out_channels=self.out_channels,
+                kernel_size=kernel_size,
+                stride=self.stride,
+                padding=padding,
+                groups=self.groups,
+                bias=False,
+            ),
+        )
+        mod_list.add_module("bn", nn.BatchNorm2d(num_features=self.out_channels))
+        return mod_list
+
+
+def reparameterize_model(model: torch.nn.Module) -> nn.Module:
+    """Method returns a model where a multi-branched structure
+        used in training is re-parameterized into a single branch
+        for inference.
+
+    Args:
+        model: MobileOne model in train mode.
+
+    Returns:
+        MobileOne model in inference mode.
+    """
+    # Avoid editing original graph
+    model = copy.deepcopy(model)
+    for module in model.modules():
+        if hasattr(module, "reparameterize"):
+            module.reparameterize()
+    return model