Export Segment Anything Model to C++

William: <p>Hi everyone,<br>
I’m currently trying to export Sam (Segment Anything model) with pytorch to use it in C++. However, it seems that it doesn’t succeed managing other classes. I saw many topics but I’m still unable to understand why it couldn’t recognize Sam type.</p>
<p>what I’m trying to do:</p>
<pre><code class="lang-auto">from segment_anything import SamPredictor, sam_model_registry
import torch


sam = sam_model_registry["default"](checkpoint="sam_vit_h_4b8939.pth")
sam.to("cuda")
predictor = SamPredictor(sam)

sm = torch.jit.script(predictor)
sm.save("sam.pt")
</code></pre>
<p>The error :</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "my_torch_project_folder_path\from_py_to_cpp.py", line 9, in &lt;module&gt;
    sm = torch.jit.script(predictor)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "my_torch_project_folder_path\.venv\Lib\site-packages\torch\jit\_script.py", line 1405, in script
    return torch.jit._recursive.create_script_class(obj)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "my_torch_project_folder_path\.venv\Lib\site-packages\torch\jit\_recursive.py", line 525, in create_script_class
    _compile_and_register_class(type(obj), rcb, qualified_class_name)
  File "my_torch_project_folder_path\.venv\Lib\site-packages\torch\jit\_recursive.py", line 61, in _compile_and_register_class
    script_class = torch._C._jit_script_class_compile(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError:
Unknown type name 'Sam':
  File "my_torch_project_folder_path\.venv\Lib\site-packages\segment_anything\predictor.py", line 20
    def __init__(
        self,
        sam_model: Sam,
                   ~~~ &lt;--- HERE
    ) -&gt; None:
        """
</code></pre>
<p>The constructor in the class</p>
<pre><code class="lang-auto">import numpy as np
import torch

from segment_anything.modeling import Sam

from typing import Optional, Tuple

from .utils.transforms import ResizeLongestSide


class SamPredictor:
    def __init__(
        self,
        sam_model: Sam,
    ) -&gt; None:
        """
        Uses SAM to calculate the image embedding for an image, and then
        allow repeated, efficient mask prediction given prompts.

        Arguments:
          sam_model (Sam): The model to use for mask prediction.
        """
        super().__init__()
        self.model = sam_model
        self.transform = ResizeLongestSide(sam_model.image_encoder.img_size)
        self.reset_image()
</code></pre>
<p>Does anyone have any explanation or solution for me? Do I have to export to the onnx format and use onnxruntime library to do this kind of stuff ? torch.jit.script seemed to be a pretty good solution to what I wanted to do. I feel like I’m missing out on something important and it is hard for me to give up.</p>
<p>Thank you very much for your help.</p>

William: <p>Just after I made this post, I looked again at my old research and went through this answer: <a href="https://stackoverflow.com/a/62647857" class="inline-onebox" rel="noopener nofollow ugc">Why can't I use classes for type annotation of the arguments to a function under the `torch.jit.script` decorator? - Stack Overflow</a><br>
It’s seems that it’s not possible to use his own class than those that are written there : <a href="https://pytorch.org/docs/stable/jit_language_reference.html#supported-type" class="inline-onebox" rel="noopener nofollow ugc">TorchScript Language Reference — PyTorch 2.3 documentation</a></p>

yluu: <p>Hi William, did you get SAM working in C++?</p>