Add the Segment Anything Model to KerasCV #1987

tirthasheshpatel · 2023-07-28T22:24:34Z

What does this PR do?

This PR implements the Segment Anything Model in multi-backend Keras.

Fixes #1679
See also #1933

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue? Please add a link
to it if that's the case.
Did you write any new necessary tests?
If this adds a new model, can you run a few training steps on TPU in Colab to ensure that no XLA incompatible OP are used?

Who can review?

@ianstenbit @DavidLandup0

ianstenbit

Thanks for the PR!

This is super exciting 🎊

Just a few comments as I took a quick look

ianstenbit · 2023-07-30T11:58:10Z

keras_cv/models/segmentation/segment_anything/sam_layers.py

+
+
+@keras.utils.register_keras_serializable(package="keras_cv")
+class MLPBlock(keras.layers.Layer):


It seems a bit heavy to make this a class since we can just make this a pair of dense layers in a sequential wherever it's used.

It is used a few times in the mask decoder actually. So, just inline the dense layers would just duplicate a lot of code. Is there any side effect of having this? If not, I'd prefer keeping it.

If it's re-used in many places then I am alright with it -- it looked to me like it was only used once or twice but I probably missed some uses

ianstenbit · 2023-07-30T11:58:36Z

keras_cv/models/segmentation/segment_anything/sam_layers.py

+
+
+@keras.utils.register_keras_serializable(package="keras_cv")
+class SAMLayerNormalization(keras.layers.Layer):


Is there no way to parameterize keras.layers.LayerNormalization to achieve this?

There is keras.layers.LayerNormalization(epsilon=1e-6). I will push this in the next batch of commits.

We should probably double check and don't take my word for it, but I'm not sure if the numerics are the same. keras.layers.LayerNormalization is:

# Compute the batch normalization. inv = 1 / ops.sqrt(variance + self.epsilon) if scale is not None: scale = ops.cast(scale, inputs.dtype) inv = inv * scale x = -mean * inv if offset is not None: offset = ops.cast(offset, inputs.dtype) x = offset + x outputs = inputs * ops.cast(inv, inputs.dtype) + ops.cast( x, inputs.dtype ) outputs = ops.cast(outputs, input_dtype) # If some components of the shape got lost due to adjustments, fix that. outputs = ops.reshape(outputs, ops.shape(inputs))

For SAM, they call it LayerNorm2d() in the official implementation, but the official impl is taken directly from Detectron2 which has BatchNorm2D and LayerNorm in turn taken from ConvNeXt: https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119

Technically, the LayerNorm re-implementation in ConvNeXt, SAM and Detectron2 shouldn't be the same LayerNorm from PyTorch.

More in these issues:

there is no need to rewrite the 'class LayerNorm(nn.Module)' facebookresearch/ConvNeXt#112

LayerNorm, what is going on? facebookresearch/ConvNeXt#136

Thanks for linking these issues @DavidLandup0, I too was wondering why this was reimplemented. After some testing, I can confirm that keras.layers.LayerNormalization(epsilon=1e-6) is numerically equivalent to SAMLayerNormalization() for segment anything. Here's the code I used to test:

import os os.environ["KERAS_BACKEND"] = "torch" import numpy as np import torch import keras_core as keras from keras_cv.models.segmentation.segment_anything import sam_layers sam_ln = sam_layers.SAMLayerNormalization() ln = keras.layers.LayerNormalization(epsilon=1e-6) sam_ln.build((1, 512, 512, 3)) ln.build((1, 512, 512, 3)) sam_ln.set_weights(ln.weights) x_np = np.random.randint(0, 256, size=(1, 512, 512, 3), dtype=np.uint8) x_np = x_np.astype(np.float32) x = torch.tensor(x_np, requires_grad=True) x_sam = torch.tensor(x_np, requires_grad=True) x_out_sam = sam_ln(x_sam) x_out = ln(x) x_out_sam.backward(torch.ones_like(x_out_sam)) x_out.backward(torch.ones_like(x_out)) np.testing.assert_allclose( x_out_sam.detach().numpy(), x_out.detach().numpy(), rtol=8e-5 ) np.testing.assert_allclose( ln.weights[0].value.grad.detach().numpy(), sam_ln.weights[2].value.grad.detach().numpy(), rtol=6e-7 ) np.testing.assert_allclose( ln.weights[1].value.grad.detach().numpy(), sam_ln.weights[3].value.grad.detach().numpy() ) np.testing.assert_allclose( x_sam.grad.detach().numpy(), x.grad.detach().numpy(), atol=3e-7 )

Awesome, thanks for checking! That simplifies things a lot :D

ianstenbit · 2023-07-30T12:02:36Z

keras_cv/models/segmentation/segment_anything/sam_mask_decoder.py

+        image_pe,
+        sparse_prompt_embeddings,
+        dense_prompt_embeddings,
+        multimask_output,


should this have a default?

I think it should. I haven't set any sensible defaults yet. I will update all the layers with some default values that make sense.

ianstenbit · 2023-07-30T12:08:12Z

keras_cv/models/segmentation/segment_anything/sam.py

@@ -0,0 +1,13 @@
+# Copyright 2023 The KerasCV Authors


I think we should probably have a top-level SegmentAnything model which takes an ImageEncoder as a backbone and subclasses Task. Then the high-level workflows can live on that model.

Then we can also include a preset which includes your ported weights!

Let's also include a reference to the paper and original implementation

The image encoder for SAM is a near 1:1 from Detectron2's ViTDet.
IMO it makes sense to have ViTDet as a standalone class/network rather than a SAM encoder only.
That way we can train it from scratch, have it as a standalone object detection model, a backbone for SAM and reuse the same code across all of that

I agree @DavidLandup0, I will move the layer to keras_cv/layers.

I think we should probably have a top-level SegmentAnything model which takes an ImageEncoder as a backbone and subclasses Task. Then the high-level workflows can live on that model.

Then we can also include a preset which includes your ported weights!

That's the plan. I will add a Task model in the kera_cv/models/segmentation/segment_anything/sam.py file. I am not yet sure how exactly the training step would be implemented with the Task API, we could just raise a NotImplementedError for now and write a train step as a follow-up.

Update: I have moved the image encoder to a standalone backbone. Let me know if that looks good to you. Thanks for the suggestion @DavidLandup0!

(I will add a Task model next)

ianstenbit

Great progress!

ianstenbit · 2023-08-02T14:41:06Z

keras_cv/layers/detectron2_layers.py

+from keras_cv.models.segmentation.segment_anything.sam_layers import MLPBlock
+
+
+def get_rel_pos(query_size, key_size, rel_pos):


Let's move the helper functions to the bottom of the file

Addressed in ac7f30e

ianstenbit · 2023-08-02T14:41:47Z

keras_cv/layers/detectron2_layers_test.py

+        x_out = ops.convert_to_numpy(attention_with_rel_pe(x))
+        self.assertEqual(x_out.shape, (1, 64, 64, 1280))
+
+    def test_windowed_transformer_encoder(self):


Let's add a test for ViTDetPatchingAndEmbedding as well

Addressed in ac7f30e

ianstenbit · 2023-08-02T14:43:10Z

keras_cv/models/backbones/detectron2/detectron2_backbone.py

+from keras_cv.utils.python_utils import classproperty
+
+
+@keras.utils.register_keras_serializable(package="keras_cv.models")


Sorry that this changed while this PR is in-flight, but if you sync to master this should now be

from keras_cv.api_export import keras_cv_export @keras_cv_export("keras_cv.models.ViTDetBackbone")

(Same for all new public API symbols)

Addressed in ac7f30e

ianstenbit · 2023-08-02T14:43:49Z

keras_cv/models/backbones/detectron2/detectron2_backbone.py

+
+    def __init__(
+        self,
+        img_size=1024,


We've standardized on input_shape for this in other backbones (accepting a tuple of (height, width, channels))

I tried to add a input_shape attribute but since I am not using the Functional syntax for building the model, it doesn't allow me to add that attribute.

In case you didn't notice, I am using the call method to specify the computations in the model instead of passing a symbolic input through each layer/operation in the __init__ method.

I noticed that it's just easier not to deal with symbolic inputs with Keras Core. One of the main reasons why Keras Core struggles with symbolic inputs is that it doesn't do shape inference. For example, in Keras, this works:

from tensorflow import keras x = keras.Input([2, 3]) tf.shape(x)[0] * 10 # Note that even though the shape at axis 0 is None, # TensorFlow returns a symbolic tensor making computations # like these valid instead of throwing an exception.

but Keras Core fails, since x.shape[0] is None.

I think it should not be difficult to convert the implementation to fully use the Functional syntax but we would have to manually check if the shapes we are getting are Nones or not. So, for example, we could do this in Keras Core:

import keras_core from keras_core import ops x = keras_core.Input([2, 3]) if x.shape[0] is not None: x = ops.reshape(x, (x.shape[0] * 2, 3)) else: x = ops.reshape(x, (None, 3))

or use some other operation.

Sorry for not highlighting these details properly beforehand! I will add a bunch of comments where I do something unintuitive so it's easier for future reviewers.

OK, turns out the shape errors were not a problem here. So, got the model ported to the Functional syntax and it should now be consistent with other backbones. I have added include_rescaling, and input_tensor arguments along with the input_shape arguments. I also tested that the weights port and can be saved/loaded in any backend. Let me know if this resolves the consistency issues.

ianstenbit · 2023-08-02T14:45:25Z

keras_cv/models/backbones/detectron2/detectron2_backbone.py

+
+
+@keras.utils.register_keras_serializable(package="keras_cv.models")
+class ViTDetBackbone(Backbone):


I assume that this backbone doesn't produce pyramid_level_outputs since it's a transformer architecture -- let's call this out in the docstring, and maybe even create an @Property for self.pyramid_level_inputs which throws a nice NotImplementedError

Addressed in ac7f30e

The backbone outputs the same shapes itself, but they do use a feature pyramid output: https://arxiv.org/pdf/2203.16527.pdf

TL:DR for the paper, the simple feature pyramid on the right turned out to be the most performant for them.

Hi @DavidLandup0, thanks for pointing out the paper! I didn't know the authors also proposed a FPN! I did look into the paper but don't know how the pyramid-level inputs would fit in the backbone here. Given this PR has already blown up a bit, I'd prefer to do this as a follow-up. Maybe you can take it up if you have time :)

Doing this as a follow-up sgtm

ianstenbit · 2023-08-02T14:46:44Z

keras_cv/models/segmentation/segment_anything/sam_mask_decoder.py

+
+
+@keras.utils.register_keras_serializable(package="keras_cv")
+class MLP(keras.layers.Layer):


Is this different from MLPBlock?

We could unite those. The only difference is that the MLPBlock has architecture embedding_dim -> mlp_dim -> embedding_dim while MLP has architecture input_dim -> [hidden_dim] * (num_layers - 1) -> output_dim. Looks like a low-hanging fruit, will address in the next commit.

ianstenbit · 2023-08-02T14:47:07Z

keras_cv/models/segmentation/segment_anything/sam_test.py

+from keras_cv.tests.test_case import TestCase
+
+
+class TestSAM(TestCase):


nit: SAMTest

ianstenbit · 2023-08-15T18:43:03Z

@tirthasheshpatel LMK when you're ready for another review on this 😄

They both behave exactly the same when moving_mean and moving_variance are None and epsilon is 1e-6

…ckend

- Use `keras_cv.export_api.keras_cv_export` instead of `keras.saving.register_keras_serializable`. - Add a `SerializableSequential` class to address the saving bug with the `Sequential` model. - Push the helper functions in `keras_cv/layers/detectron2_layers.py` to the bottom of the file. - Add the detectron2 layers to the `keras_cv/layers/__init__.py` file. - Add a test for the `ViTDetPatchingAndEmbedding` layer.

keras_cv/models/backbones/detectron2/detectron2_backbone.py

tirthasheshpatel · 2023-08-21T08:38:14Z

Hi @ianstenbit, thank you very much for your reviews so far! Very helpful!

LMK when you're ready for another review on this

I have a Task API for the SAM model ready, I will just have to make some more changes before pushing it + add tests. I will mark the PR ready for review once I do that and help start reviews. But if you have time, feel free to review!

DavidLandup0 · 2023-08-21T17:21:09Z

keras_cv/layers/__init__.py

@@ -17,6 +17,10 @@
 from tensorflow.keras.layers import RandomWidth

 from keras_cv.layers.augmenter import Augmenter
+from keras_cv.layers.detectron2_layers import AddPositionalEmbedding


Do we want to put these under a detectron2 namespace?

DavidLandup0 · 2023-08-21T17:22:57Z

keras_cv/layers/__init__.py

@@ -17,6 +17,10 @@
 from tensorflow.keras.layers import RandomWidth

 from keras_cv.layers.augmenter import Augmenter
+from keras_cv.layers.detectron2_layers import AddPositionalEmbedding


Since this would be exported as part of the public API - we have PatchingAndEmbedding which does patching with a Conv2D and then adds embeddings in this same form. Do we want to update that layer to use AddPositionalEmbedding as well for conformity?

Wouldn't adding a new layer in a pre-existing class invalidate the weights set for the ViT model? Also, since PatchingAndEmbedding is still a TensorFlow Keras layer, I think, for the time being, it'd be easier to keep the two separate.

Although, I'd add a comment about this as a TODO so we don't forget to do it in the future. What do you think?

DavidLandup0 · 2023-08-21T17:24:55Z

keras_cv/layers/detectron2_layers.py

+
+
+@keras_cv_export("keras_cv.layers.MultiHeadAttentionWithRelativePE")
+class MultiHeadAttentionWithRelativePE(keras.layers.Layer):


Perhaps it would make sense to do an AddRelativePositionalEmbedding class for consistency with the aforementioned AddPositionalEmbedding?

DavidLandup0 · 2023-08-21T17:26:20Z

keras_cv/layers/detectron2_layers.py

+        )
+
+        if self.use_rel_pos:
+            attention_map = add_decomposed_rel_pos(


This should probably be a private method as part of a layer

DavidLandup0 · 2023-08-21T17:27:37Z

keras_cv/layers/detectron2_layers.py

+        if self.window_size > 0:
+            H, W = x.shape[1], x.shape[2]
+
+            x, HW_padded = window_partition(x, self.window_size)


What do you think about doing this as a layer instead of a method?

I.e. https://github.com/DavidLandup0/deepvision/blob/main/deepvision/layers/window_partitioning.py

Done. Instead of creating two classes, one for partitioning and one for unpartitioning, I handled both in a single class. Let me know if that looks good.

DavidLandup0 · 2023-08-21T17:29:16Z

keras_cv/layers/detectron2_layers.py

+
+
+@keras_cv_export("keras_cv.layers.ViTDetPatchingAndEmbedding")
+class ViTDetPatchingAndEmbedding(keras.layers.Layer):


This is the same as the ViT patching and embedding but without positional embedding.
I'm torn between being able to turn off PE in the default layer and adding that as a flag and having a new layer for this...

Are you referring to the PathingAndEmbedding class for the ViT model in KerasCV? I addressed that here: #1987 (comment)

DavidLandup0 · 2023-08-21T17:30:09Z

keras_cv/layers/detectron2_layers.py

+        return config
+
+
+def get_rel_pos(query_size, key_size, rel_pos):


This should probably be a private method or turned into a public layer.
I.e.: https://github.com/DavidLandup0/deepvision/blob/main/deepvision/layers/decomposed_relative_positional_embedding.py

DavidLandup0 · 2023-08-21T17:30:19Z

keras_cv/layers/detectron2_layers.py

+    return ops.take(rel_pos_resized, relative_coordinates, 0)
+
+
+def add_decomposed_rel_pos(


Same comment as above

DavidLandup0 · 2023-08-21T17:30:47Z

keras_cv/layers/serializable_sequential.py

+# This only happens when the `build` method is called in the `__init__`
+# step.
+@keras_cv_export("keras_cv.layers.SerializableSequential")
+class SerializableSequential(keras.layers.Layer):


Is this still an issue in Keras Core?

The bug has been addressed in Keras Core v0.1.5 but the latest TensorFlow Keras still has it. So, weights won't load in TensorFlow Keras until the bug is addressed in the next release.

We can either:

Drop support temporarily for TensorFlow Keras just for this model with a note in the docs. With the new release of TF Keras, the bug should be fixed and we can remove the note.

Keep the simple replication of the class until the bug is resolved in TensorFlow Keras.

I am leaning more towards option 2 but I don't have a strong opinion. What do you think @ianstenbit @DavidLandup0?

I ended up removing it since the legacy weights load in all backends in Keras Core and also in TF Keras. I think until some saving issues are addressed with the new .weights.h5 format, we should just use the legacy weights. Let me know what you both think!

DavidLandup0 · 2023-08-21T17:32:02Z

keras_cv/models/__init__.py

@@ -43,6 +43,18 @@
 from keras_cv.models.backbones.densenet.densenet_backbone import (
    DenseNetBackbone,
 )
+from keras_cv.models.backbones.detectron2.detectron2_aliases import (


Afaik, the backbone is basically the same as the official ViTDet, so there may not be a need to call it a SAM{name}Backbone

Sounds good, thanks for looking into it!

DavidLandup0 · 2023-08-21T17:32:25Z

keras_cv/models/__init__.py

@@ -166,5 +178,8 @@
    YOLOV8Detector,
 )
 from keras_cv.models.segmentation import DeepLabV3Plus
+from keras_cv.models.segmentation import MaskDecoder


May be better as SAMMaskDecoder for clarity

DavidLandup0 · 2023-08-21T17:32:48Z

keras_cv/models/backbones/detectron2/data/sam_vitdet_b_out.npz

Probably left by accident?

I have added this intentionally. This is used in the tests to verify that the model weights are loaded correctly and that the forward pass in all backends yields the same result.

DavidLandup0 · 2023-08-21T17:36:39Z

keras_cv/models/backbones/detectron2/detectron2_backbone.py

+        """Dictionary of preset names and configurations."""
+        return copy.deepcopy(backbone_presets)
+
+    # @classproperty


stray comment?

This method loads the presets with weights. I will uncomment it later once the model layers are finalized and the final weights are uploaded.

DavidLandup0 · 2023-08-21T17:38:00Z

keras_cv/models/segmentation/segment_anything/sam_layers.py

+
+
+@keras_cv_export("keras_cv.layers.MLP")
+class MLP(keras.layers.Layer):


This is a public class here - it should probably be a private subclass, especially since there was an MLP with the same name in a layer related to this, iirc

Done. I have removed the export.

DavidLandup0 · 2023-08-21T17:38:47Z

keras_cv/models/segmentation/segment_anything/sam_mask_decoder.py

@@ -0,0 +1,230 @@
+# Copyright 2023 The KerasCV Authors


I'd probably put this layer under sam layers

DavidLandup0 · 2023-08-21T17:39:01Z

keras_cv/models/segmentation/segment_anything/sam_mask_decoder.py

+
+@keras_cv_export("keras_cv.models.MaskDecoder")
+class MaskDecoder(keras.models.Model):
+    """Mask decoder for the segment anything model.


Nit: "Segment Anything (SAM)"

DavidLandup0 · 2023-08-21T17:39:22Z

keras_cv/models/segmentation/segment_anything/sam_mask_decoder.py

+
+
+@keras_cv_export("keras_cv.models.MaskDecoder")
+class MaskDecoder(keras.models.Model):


As mentioned before, to avoid confusion, probably best if this is called SAMMaskDecoder or something along those lines

DavidLandup0 · 2023-08-21T17:39:34Z

keras_cv/models/segmentation/segment_anything/sam_mask_decoder.py

+            network. Defaults to "gelu".
+
+    References:
+        - [Segment Anything](https://arxiv.org/abs/2304.02643)


We probably want a code reference as well

[skip ci]

tirthasheshpatel · 2023-09-11T08:25:55Z

I think the PR is almost ready for some thorough reviews except for a few TODOs:

I ported the weights to the Keras Core model and am able to load in any backend but loading weights is broken between Keras Core and TensorFlow Keras (xref Saving broken between tf.keras and Keras Core keras-core#855)
I need to add more docs (especially examples) for the internal layers that are exported.
Need to upload weights and add them as presets here. I think we need to first finalize a few blockers (point 1 above and this discussion) before uploading the final weights set.

Let me know if you have any other major points @DavidLandup0 @ianstenbit. And thanks for the reviews so far, super helpful!

tirthasheshpatel · 2023-09-14T20:10:33Z

I ported the weights to the Keras Core model and am able to load in any backend but loading weights is broken between Keras Core and TensorFlow Keras (xref

An update on this: the legacy weights *.h5 load in both TF Keras and Keras Core (all backends). I think we can just use that until the saving fixes are available in TF Keras. Also, SerializableSequential is no longer needed when using legacy weights.

ianstenbit

The only thing left that I see is adding a preset for the pre-trained version of SAM.

Thanks for your great work!

keras_cv/models/backbones/detectron2/detectron2_backbone.py

ianstenbit · 2023-09-15T16:54:55Z

keras_cv/models/backbones/detectron2/detectron2_backbone.py

+
+
+@keras.utils.register_keras_serializable(package="keras_cv.models")
+class ViTDetBackbone(Backbone):


Doing this as a follow-up sgtm

keras_cv/models/segmentation/segment_anything/sam.py

ianstenbit · 2023-09-18T15:47:12Z

/gcbrun

ianstenbit

This is awesome -- thanks Tirth!

Just one little fix to make GCBRun happy

keras_cv/models/segmentation/segment_anything/sam_test.py

ianstenbit · 2023-09-18T20:04:25Z

/gcbrun

tirthasheshpatel · 2023-09-19T03:18:22Z

Thanks, @DavidLandup0 @ianstenbit for your reviews! This was fun to work on. Excited to have this in KerasCV!

The next steps are to add some guides to use and train the model. It would also be nice to have some benchmarks. On it now! But I will also create a tracking issue in case the community wants to take over some of these tasks.

ianstenbit · 2023-09-19T03:19:55Z

Thanks, @DavidLandup0 @ianstenbit for your reviews! This was fun to work on. Excited to have this in KerasCV!

The next steps are to add some guides to use and train the model. It would also be nice to have some benchmarks. On it now! But I will also create a tracking issue in case the community wants to take over some of these tasks.

Thank you Tirth for your outstanding work on this -- we really appreciate it!

I think our long-term goal should be to add support for text prompts. There are some community projects out there which demonstrate the feasibility of this, and I think it would be a great step for us.

But I 100% agree that some guides and training are the right place to start!

* Start adding components for the segment anything model * SAMLayerNormalization -> keras.layers.LayerNormalization They both behave exactly the same when moving_mean and moving_variance are None and epsilon is 1e-6 * Move the image encoder to detectron2 backbone and fix for tf.keras backend * Address review comments and address saving bug - Use `keras_cv.export_api.keras_cv_export` instead of `keras.saving.register_keras_serializable`. - Add a `SerializableSequential` class to address the saving bug with the `Sequential` model. - Push the helper functions in `keras_cv/layers/detectron2_layers.py` to the bottom of the file. - Add the detectron2 layers to the `keras_cv/layers/__init__.py` file. - Add a test for the `ViTDetPatchingAndEmbedding` layer. * Make the backbone functional; unite MLP and MLPBlock * Address David's review comments * Add SAM Task model; make MaskDecoder and PromptEncoder XLA compatible * Remove a stray file * Add docs for the Task model * Add more references [skip ci] * Remove SerializableSequential layer * detectron2 -> vit_det; add SAM presets; fix ViTDet presets * Increse test tolerence for GCB Run

ianstenbit reviewed Jul 30, 2023

View reviewed changes

ianstenbit reviewed Aug 2, 2023

View reviewed changes

tirthasheshpatel added 3 commits August 18, 2023 11:07

Start adding components for the segment anything model

ddebecd

SAMLayerNormalization -> keras.layers.LayerNormalization

025dd15

They both behave exactly the same when moving_mean and moving_variance are None and epsilon is 1e-6

Move the image encoder to detectron2 backbone and fix for tf.keras ba…

43b0f2b

…ckend

tirthasheshpatel force-pushed the add-sam branch from 17a1cbb to 05d1d27 Compare August 18, 2023 21:45

tirthasheshpatel force-pushed the add-sam branch from 05d1d27 to ac7f30e Compare August 18, 2023 21:47

Make the backbone functional; unite MLP and MLPBlock

360baf7

tirthasheshpatel commented Aug 21, 2023

View reviewed changes

keras_cv/models/backbones/detectron2/detectron2_backbone.py Outdated Show resolved Hide resolved

DavidLandup0 reviewed Aug 21, 2023

View reviewed changes

tirthasheshpatel added 2 commits August 24, 2023 23:29

Address David's review comments

d7a1eb3

Add SAM Task model; make MaskDecoder and PromptEncoder XLA compatible

ada6040

tirthasheshpatel marked this pull request as ready for review September 9, 2023 23:13

tirthasheshpatel changed the title ~~[WIP] Add the Segment Anything Model to KerasCV~~ Add the Segment Anything Model to KerasCV Sep 9, 2023

tirthasheshpatel requested review from ianstenbit and DavidLandup0 September 9, 2023 23:28

tirthasheshpatel added 4 commits September 9, 2023 16:33

Merge branch 'master' of github.com:keras-team/keras-cv into add-sam

93eb78f

Remove a stray file

d908527

Add docs for the Task model

07cb9d6

Add more references

343557a

[skip ci]

divyashreepathihalli mentioned this pull request Sep 14, 2023

Segment Anything Model inside Keras Applications keras-team/keras#18407

Closed

Remove SerializableSequential layer

19312e7

ianstenbit reviewed Sep 15, 2023

View reviewed changes

detectron2 -> vit_det; add SAM presets; fix ViTDet presets

e5767d8

ianstenbit approved these changes Sep 18, 2023

View reviewed changes

keras_cv/models/segmentation/segment_anything/sam_test.py Outdated Show resolved Hide resolved

tirthasheshpatel commented Sep 18, 2023

View reviewed changes

keras_cv/models/segmentation/segment_anything/sam_test.py Outdated Show resolved Hide resolved

Increse test tolerence for GCB Run

8ad94ae

ianstenbit merged commit bc80fbb into keras-team:master Sep 19, 2023
8 of 9 checks passed

tirthasheshpatel mentioned this pull request Sep 19, 2023

Segment Anything - Next Steps #2081

Open

5 tasks

tirthasheshpatel deleted the add-sam branch September 19, 2023 03:33

tirthasheshpatel restored the add-sam branch September 19, 2023 06:29

tirthasheshpatel deleted the add-sam branch September 19, 2023 06:29



		@keras.utils.register_keras_serializable(package="keras_cv")
		class MLPBlock(keras.layers.Layer):



		@keras.utils.register_keras_serializable(package="keras_cv")
		class SAMLayerNormalization(keras.layers.Layer):

		from keras_cv.models.segmentation.segment_anything.sam_layers import MLPBlock


		def get_rel_pos(query_size, key_size, rel_pos):

		from keras_cv.utils.python_utils import classproperty


		@keras.utils.register_keras_serializable(package="keras_cv.models")



		@keras.utils.register_keras_serializable(package="keras_cv.models")
		class ViTDetBackbone(Backbone):



		@keras.utils.register_keras_serializable(package="keras_cv")
		class MLP(keras.layers.Layer):

		from keras_cv.tests.test_case import TestCase


		class TestSAM(TestCase):



		@keras_cv_export("keras_cv.layers.MultiHeadAttentionWithRelativePE")
		class MultiHeadAttentionWithRelativePE(keras.layers.Layer):



		@keras_cv_export("keras_cv.layers.ViTDetPatchingAndEmbedding")
		class ViTDetPatchingAndEmbedding(keras.layers.Layer):

		return config


		def get_rel_pos(query_size, key_size, rel_pos):

		return ops.take(rel_pos_resized, relative_coordinates, 0)


		def add_decomposed_rel_pos(



		@keras_cv_export("keras_cv.layers.MLP")
		class MLP(keras.layers.Layer):



		@keras_cv_export("keras_cv.models.MaskDecoder")
		class MaskDecoder(keras.models.Model):

Add the Segment Anything Model to KerasCV #1987

Add the Segment Anything Model to KerasCV #1987

Conversation

tirthasheshpatel commented Jul 28, 2023 • edited Loading

What does this PR do?

Before submitting

Who can review?

ianstenbit left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tirthasheshpatel Jul 31, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tirthasheshpatel Jul 31, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tirthasheshpatel Jul 31, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tirthasheshpatel Jul 31, 2023 • edited Loading

Choose a reason for hiding this comment

tirthasheshpatel Aug 2, 2023 • edited Loading

Choose a reason for hiding this comment

ianstenbit left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ianstenbit commented Aug 15, 2023

tirthasheshpatel commented Aug 21, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tirthasheshpatel Aug 24, 2023 • edited Loading

Choose a reason for hiding this comment

tirthasheshpatel Sep 15, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tirthasheshpatel commented Jul 28, 2023 •

edited

Loading

tirthasheshpatel Jul 31, 2023 •

edited

Loading

tirthasheshpatel Jul 31, 2023 •

edited

Loading

tirthasheshpatel Jul 31, 2023 •

edited

Loading

tirthasheshpatel Jul 31, 2023 •

edited

Loading

tirthasheshpatel Aug 2, 2023 •

edited

Loading

tirthasheshpatel Aug 24, 2023 •

edited

Loading

tirthasheshpatel Sep 15, 2023 •

edited

Loading