Running facebook's segmentation model using ortex

hmm sorry I thought i had one but I sent the wrong link, still looking for the right one

update: vietanhdev/segment-anything-2-onnx-models at main

yes there was indeed an error on the gist sorry. However the livebook that @kip made is fully functional.

I’ll try to make it work with SAM 2 but I have yet to fin ONNX models of it.

I tried to run the livebook from here
https://raw.githubusercontent.com/elixir-image/image/main/livebook/segment_anything.livemd

It gives me
** (ArgumentError) cannot broadcast tensor of dimensions {3, 768, 1024} to {1, 3, 1024, 1024} with axes [1, 2, 3]
(nx 0.7.3) lib/nx/shape.ex:253: Nx.Shape.broadcast!/4

I had to change the resizing line to

resized =
  Image.thumbnail!(image, "1024x1024", crop: :attention)

to make it end up as square. Other options are here Image — image v0.48.1

Would love to have this work with SAM 2

That’s amazing, thank you!
Indeed, the image library makes things much simpler and cleaner. I didn’t know these badges existed, awesome!

Happy to help! I also learned a lot to get the code working.

I’ve made a modified version of your Livebook to use Image for the image pipeline. I think it simplifies some of the code.

I’ve also made it so the encoder and decoder are downloaded from hugging face using req so the Livebook now works standalone. So here’s the badge!

Thanks for taking the time to help me through my lack of understanding. Its been a good learning experience.

Really appreciate it, thanks very much. Back to experimental mode!

That’s my bad, with all my different trials I got mixed up in my code versions…

I’ve uploaded my complete working code and models on hugging face which you can find here : ginkgoo/SegmentAnythingModel-Elixir-Ortex at main

That should do the work !

Seems I’m still a bit stuck - and hoping your patience hasn’t run out

Using the encoder link you pointed me at, I’m seeing the following error:

tensor =
  image_tensor
  |> Nx.as_type(:f32)

# Mean and std values copied from transformer.js
mean = Nx.tensor([123.675, 116.28, 103.53])
std = Nx.tensor([58.395, 57.12, 57.375])

normalized_tensor =
  tensor
  |> NxImage.normalize(mean, std)

{image_embeddings} = Ortex.run(encoder, Nx.broadcast(normalized_tensor, {1024, 1024, 3}))

And the error is:

** (RuntimeError) Dimensions do not match: InputsLength { inference_input: [[1024, 1024, 3]], model_input: [[None, Some(3), Some(1024), Some(1024)]] }
    (ortex 0.1.9) lib/ortex/model.ex:51: Ortex.Model.run/2
    #cell:le3cboucvteyzzeh:21: (file)

Which I think is saying that the shape of the image data is {1024, 1024, 3} but the model expects {3, 1024, 1024}.

Given you already have a working model, may I still ask if you are able to upload the encoder and decoder you are using so I can at least eliminate that as an issue in my experiment?

@herisson, much appreciated - I will give that a go and see where I get to. Many thanks.

Yes, their examples aren’t very comprehensive in terms of explanations. From what I understand, SAM is a “two-stage model.”

The first stage takes an image and transforms it into something the next model can understand (image embeddings). This stage, which I refer to as the encoder or vision encoder, is common in many image processing tasks, not unique to SAM. That’s why it’s not included in their export script.

The second stage, which you exported, takes the image embeddings and other inputs to produce the mask (referred to as the decoder).

I can’t upload the decoder to a repository right now as I’m not at home, but you can find the vision encoder here. Since you’ve exported the decoder, you should be set with these two models!

Hope it makes things clearer