sam2

package
v0.4.0 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Jul 21, 2026 License: Apache-2.0 Imports: 26 Imported by: 0

README

SAM2: Segment Anything Model 2 in GoMLX

This package implements the image segmentation part of Facebook's Segment Anything in Images And Videos (SAM2) foundation model. It is designed to run efficiently in Go using GoMLX for hardware-accelerated tensor computations.

This package implements the image part of the model (including the Hiera vision backbone, neck, prompt encoder, and mask decoder), not the video tracking part.

References

APIs Offered

The package exposes two levels of APIs:

1. High-Level Go API (Inference Segmenter)

A user-friendly, standard Go interface that operates on image.Image inputs and handles raw tensor conversion, preprocessing, and output resizing automatically.

  • NewSegmenter: Initializes a predictor using the loaded config/weights and compiles the computation graphs.
  • Segment: Predicts masks for a given image and prompt options.
2. Low-Level Graph API (GoMLX Graph Construction)

A graph-building API designed for users who want to embed SAM2 inside their custom GoMLX computational graphs, training loops, or pipelines.

  • Forward: Constructs the full SAM2 graph (Vision Encoder $\rightarrow$ Prompt Encoder $\rightarrow$ Mask Decoder) using GoMLX *Node operations.
  • InternalStates: A struct exposing intermediate outputs (e.g. FPN neck features, prompt embeddings, upscaled mask maps) for transfer learning and debugging.

Code Structure

  • sam2.go: The core model architecture (backbone, neck, attention layers, and mask decoder).
  • goapi.go: The standard Go inference API wrapping the GoMLX executable.
  • model.go: SafeTensors weights loading, variable mapping, and transpositions.
  • config.go: Configuration parsing and model parameters.

Usage Examples

Using the High-Level Segmenter
package main

import (
	"fmt"
	"image"
	_ "image/png"
	"log"
	"os"

	"github.com/gomlx/compute"
	"github.com/gomlx/go-huggingface/hub"
	"github.com/gomlx/go-huggingface/models/sam2"
	_ "github.com/gomlx/gomlx/backends/default"
)

func main() {
	// Initialize GoMLX backend (e.g., CUDA, CPU)
	backend, err := compute.New()
	if err != nil {
		log.Fatalf("Failed to initialize backend: %v", err)
	}
	defer backend.Finalize()

	// Load configuration and weights from HuggingFace
	repo := hub.New("facebook/sam2-hiera-base-plus")
	modelObj, err := sam2.LoadModel(repo)
	if err != nil {
		log.Fatalf("Failed to load model: %v", err)
	}

	// Create standard Segmenter
	segmenter, err := sam2.NewSegmenter(backend, modelObj)
	if err != nil {
		log.Fatalf("Failed to create segmenter: %v", err)
	}

	// Load target image
	imgFile, err := os.Open("target.png")
	if err != nil {
		log.Fatalf("Failed to open image: %v", err)
	}
	defer imgFile.Close()
	img, _, err := image.Decode(imgFile)

	// Set point prompt (foreground point at x=512, y=512)
	options := &sam2.PredictOptions{
		Points: []sam2.PromptPoint{
			{X: 512, Y: 512, Label: sam2.LabelForeground},
		},
		MultiMaskOutput: false,
	}

	// Segment image
	segmentations, err := segmenter.Segment(img, options)
	if err != nil {
		log.Fatalf("Segmentation failed: %v", err)
	}

	best := segmentations[0]
	fmt.Printf("Predicted mask with IoU score: %.4f\n", best.IoUScore)
	// best.Mask contains the grey/binary mask image.Image
}

Demo CLI Program

A command-line program is available under models/sam2/demo to segment arbitrary images.

Running the Demo

  1. Build the demo binary:

    go build -o sam2-demo ./models/sam2/demo

  2. Segment an image using a point prompt (-points "x,y,label" where label 1 is foreground, 0 is background) and overlay a red highlight mask:

    ./sam2-demo -input input.png -output output.png -points "512,512,1" -color "red"

  3. Segment using a bounding box prompt (-boxes "x_min,y_min,x_max,y_max"):

    ./sam2-demo -input input.png -output output.png -boxes "100,100,800,800" -color "blue"

  4. Run in multi-mask mode to output all 3 candidate masks:

    ./sam2-demo -input input.png -output output.png -points "512,512,1" -multimask=true
Demo Flags:
  • -input: Path to the input image file (JPEG or PNG).
  • -output: Path to save the output segmented image (defaults to output.png).
  • -model: HuggingFace repository ID of the model (defaults to facebook/sam2-hiera-base-plus).
  • -points: Semicolon-separated point coordinates (e.g. x1,y1,label1;x2,y2,label2).
  • -boxes: Semicolon-separated box coordinates (e.g. xmin,ymin,xmax,ymax).
  • -color: Mask overlay highlight color (supports hex like #ff0000, RGB like 255,0,0, or names like red, green, blue, gray, yellow).
  • -multimask: Output 3 ambiguous masks as separate files (output_0.png, output_1.png, etc.).
  • -format: Output image format (png or jpg).

Documentation

Overview

Package sam2 implements the Facebook's image segmentation foundation model "Segment Anything in Images And Videos" (SAM2), see [1], [2].

This package implements the image part of the model (the Hiera backbone), not the video part. It's based on the HuggingFace implementation in [3].

In includes a Go friendly Segmenter, see NewSegmenter.

[1] https://arxiv.org/pdf/2408.00714 [2] https://huggingface.co/facebook/sam2-hiera-base-plus [3] https://github.com/huggingface/transformers/tree/main/src/transformers/models/sam2

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func BuildSinePositionEmbedding 

func BuildSinePositionEmbedding(g *Graph, batchSize, height, width int, numPositionFeatures int, normalize bool, scale float64, temperature float64, dtype dtypes.DType) *Node

BuildSinePositionEmbedding builds the Feature Pyramid Network (FPN) neck sine/cosine position embeddings.

func EmbedBoxes 

func EmbedBoxes(scope *model.Scope, boxes *Node, inputImageSize int, hiddenSize int) *Node

EmbedBoxes embeds box prompts.

func EmbedPoints 

func EmbedPoints(scope *model.Scope, points, labels *Node, pad bool, config *Config) (*Node, *Node, *Node)

EmbedPoints embeds point/label prompts.

func Forward 

func Forward(scope *model.Scope, pixelValues *Node, inputPoints, inputLabels, inputBoxes, inputMasks *Node, config *Config, multimaskOutput bool) (*Node, InternalStates)

Forward performs the full model forward pass.

Inputs:

  • scope: The model scope.
  • pixelValues: Input image tensor of shape [batch_size, 3, 1024, 1024]. It should be normalized with ImageNet mean ([0.485, 0.456, 0.406]) and std ([0.229, 0.224, 0.225]).
  • inputPoints: Input point coordinates of shape [batch_size, point_batch_size, num_points, 2]. The coordinates should be absolute pixel coordinates (e.g. from 0 to 1023). Can be nil if inputBoxes is provided.
  • inputLabels: Input point labels of shape [batch_size, point_batch_size, num_points] (values: 1 for positive, 0 for negative, -1 for padding).
  • inputBoxes: Input box coordinates of shape [batch_size, point_batch_size, 4] (format: [x_min, y_min, x_max, y_max]). Can be nil.
  • inputMasks: Input low-resolution masks of shape [batch_size, 1, 256, 256]. Can be nil.
  • config: SAM2 model configuration.
  • multimaskOutput: If true, returns 3 predicted masks per prompt. If false, returns a single predicted mask.

Outputs:

  • predictionMask: The predicted mask logit tensor of shape [batch_size, point_batch_size, num_masks, 256, 256], where num_masks is 3 if multimaskOutput is true, or 1 if false. Mask values are unnormalized logits.
  • internalStates: Struct containing intermediate states and predictions (e.g. IoUScores for Intersection over Union, FPNHiddenStates for Feature Pyramid Network) that are mainly used for testing and can be ignored for normal inference.

func GetImageWidePositionalEmbeddings 

func GetImageWidePositionalEmbeddings(scope *model.Scope, g *Graph, config *Config) *Node

GetImageWidePositionalEmbeddings builds the 2D positional embeddings for the mask decoder.

func GetPosEmbed 

func GetPosEmbed(scope *model.Scope, g *Graph, h, w int, config *Config) *Node

GetPosEmbed generates Hiera Det positional embeddings.

func RepeatInterleaveAxis0 

func RepeatInterleaveAxis0(x *Node, repeats int) *Node

RepeatInterleaveAxis0 repeats a tensor along axis 0.

func Sam2FeedForward3Layers 

func Sam2FeedForward3Layers(scope *model.Scope, x *Node, inDim, hiddenDim, outDim int, hiddenAct string) *Node

Sam2FeedForward3Layers implements a 3-layer MLP.

func Sam2HieraDetModel 

func Sam2HieraDetModel(scope *model.Scope, x *Node, config *Config) ([]*Node, *Node)

Sam2HieraDetModel runs the vision transformer backbone.

func Sam2MaskEmbedding 

func Sam2MaskEmbedding(scope *model.Scope, masks *Node, config *Config) *Node

Sam2MaskEmbedding embeds input masks.

func Sam2MultiScaleAttention 

func Sam2MultiScaleAttention(scope *model.Scope, hiddenStates *Node, dim, dimOut, numHeads int, queryStride []int) *Node

Sam2MultiScaleAttention builds the multi-head attention block used in the Hiera backbone.

func Sam2MultiScaleBlock 

func Sam2MultiScaleBlock(scope *model.Scope, x *Node, dim, dimOut, numHeads, windowSize int, queryStride []int, useGlobalAttn bool) *Node

Sam2MultiScaleBlock runs one hierarchical block (attention + FFN).

func Sam2PositionalEmbedding 

func Sam2PositionalEmbedding(scope *model.Scope, coords *Node, inputShape []int) *Node

Sam2PositionalEmbedding generates positional sine/cosine embeddings.

func Sam2VisionNeck 

func Sam2VisionNeck(scope *model.Scope, hiddenStates []*Node, config *Config) (fpnHiddenStates []*Node, fpnPositionEncoding []*Node)

Sam2VisionNeck projects backbone features and performs FPN (Feature Pyramid Network) fusion.

func WindowPartition 

func WindowPartition(x *Node, windowSize int) (windows *Node, paddedShape []int)

WindowPartition partitions input into non-overlapping windows.

func WindowUnpartition 

func WindowUnpartition(windows *Node, windowSize int, hp, wp, h, w int) *Node

WindowUnpartition merges non-overlapping windows back.

Types

type Config 

type Config struct {
	ModelType           string              `json:"model_type"`
	VisionConfig        VisionConfig        `json:"vision_config"`
	PromptEncoderConfig PromptEncoderConfig `json:"prompt_encoder_config"`
	MaskDecoderConfig   MaskDecoderConfig   `json:"mask_decoder_config"`
	TorchDtype          string              `json:"torch_dtype"`
	DType               string              `json:"dtype"`
}

Config represents the top-level configuration for SAM2.

func LoadConfig 

func LoadConfig(path string) (*Config, error)

LoadConfig parses a config.json file.

type HieraDetConfig 

type HieraDetConfig struct {
	HiddenSize                              int     `json:"hidden_size"`
	NumAttentionHeads                       int     `json:"num_attention_heads"`
	NumChannels                             int     `json:"num_channels"`
	ImageSize                               []int   `json:"image_size"`
	PatchKernelSize                         []int   `json:"patch_kernel_size"`
	PatchStride                             []int   `json:"patch_stride"`
	PatchPadding                            []int   `json:"patch_padding"`
	QueryStride                             []int   `json:"query_stride"`
	WindowPositionalEmbeddingBackgroundSize []int   `json:"window_positional_embedding_background_size"`
	NumQueryPoolStages                      int     `json:"num_query_pool_stages"`
	BlocksPerStage                          []int   `json:"blocks_per_stage"`
	EmbedDimPerStage                        []int   `json:"embed_dim_per_stage"`
	NumAttentionHeadsPerStage               []int   `json:"num_attention_heads_per_stage"`
	WindowSizePerStage                      []int   `json:"window_size_per_stage"`
	GlobalAttentionBlocks                   []int   `json:"global_attention_blocks"`
	MLPRatio                                float64 `json:"mlp_ratio"`
	HiddenAct                               string  `json:"hidden_act"`
	LayerNormEps                            float64 `json:"layer_norm_eps"`
}

HieraDetConfig represents the backbone configuration.

type InternalStates 

type InternalStates struct {
	// IoUScores predicted for each mask. IoU stands for Intersection over Union.
	// Shape: [batch_size, point_batch_size, num_masks] (values usually in [0, 1]).
	IoUScores *Node
	// FPNHiddenStates are the multi-scale features from the vision neck. FPN stands for Feature Pyramid Network.
	// It contains 3 tensors of shapes [batch_size, 32, 256, 256], [batch_size, 64, 128, 128], [batch_size, 256, 64, 64].
	FPNHiddenStates []*Node
	// SparseEmbeddings generated by the prompt encoder for points/boxes. Shape: [batch_size, point_batch_size, num_tokens, 256].
	SparseEmbeddings *Node
	// DenseEmbeddings generated by the prompt encoder (e.g. mask prompt). Shape: [batch_size, 256, 64, 64].
	DenseEmbeddings *Node
	// IoUTokenOut is the IoU (Intersection over Union) token output from the transformer. Shape: [batch_size, point_batch_size, 256].
	IoUTokenOut *Node
	// MaskTokensOut are the mask tokens output from the transformer. Shape: [batch_size, point_batch_size, 4, 256].
	MaskTokensOut *Node
	// Upscaled features from the first upscaling stage. Shape: [batch_size * point_batch_size, 64, 128, 128].
	Upscaled *Node
	// Upscaled2 features from the second upscaling stage. Shape: [batch_size * point_batch_size, 32, 256, 256].
	Upscaled2 *Node
	// HyperIn are the hypernetwork weights generated by output MLPs. Shape: [batch_size * point_batch_size, 4, 32].
	HyperIn *Node
}

InternalStates holds intermediate outputs and variables of the SAM2 model. It is primarily used for testing and debugging, but can also be used for partial transfer learning (e.g. extracting image embeddings, sparse/dense prompt embeddings, or upscaled features).

type MaskDecoderConfig 

type MaskDecoderConfig struct {
	HiddenSize                      int     `json:"hidden_size"`
	HiddenAct                       string  `json:"hidden_act"`
	MLPDim                          int     `json:"mlp_dim"`
	NumHiddenLayers                 int     `json:"num_hidden_layers"`
	NumAttentionHeads               int     `json:"num_attention_heads"`
	AttentionDownsampleRate         int     `json:"attention_downsample_rate"`
	NumMultimaskOutputs             int     `json:"num_multimask_outputs"`
	IoUHeadDepth                    int     `json:"iou_head_depth"`
	IoUHeadHiddenDim                int     `json:"iou_head_hidden_dim"`
	DynamicMultimaskViaStability    bool    `json:"dynamic_multimask_via_stability"`
	DynamicMultimaskStabilityDelta  float64 `json:"dynamic_multimask_stability_delta"`
	DynamicMultimaskStabilityThresh float64 `json:"dynamic_multimask_stability_thresh"`
}

MaskDecoderConfig represents the mask decoder configuration.

type Model 

type Model struct {
	Repo   *hub.Repo
	Config *Config
	// contains filtered or unexported fields
}

Model holds the configuration and reference to the HuggingFace repository.

func LoadModel 

func LoadModel(repo *hub.Repo) (*Model, error)

LoadModel loads the config.json and returns a Model.

func (*Model) LoadStore 

func (m *Model) LoadStore(backend compute.Backend, store *model.Store) error

LoadStore loads the safetensors weights into the GoMLX model.Store.

type PointLabel 

type PointLabel int

PointLabel indicates whether a point is foreground (positive) or background (negative). 

const (
	LabelBackground PointLabel = 0
	LabelForeground PointLabel = 1
)

type PromptBox 

type PromptBox struct {
	MinX, MinY int
	MaxX, MaxY int
}

PromptBox defines a bounding box for targeting objects.

type PromptEncoderConfig 

type PromptEncoderConfig struct {
	HiddenSize         int     `json:"hidden_size"`
	ImageSize          int     `json:"image_size"`
	PatchSize          int     `json:"patch_size"`
	MaskInputChannels  int     `json:"mask_input_channels"`
	NumPointEmbeddings int     `json:"num_point_embeddings"`
	HiddenAct          string  `json:"hidden_act"`
	LayerNormEps       float64 `json:"layer_norm_eps"`
	Scale              float64 `json:"scale"`
}

PromptEncoderConfig represents the prompt encoder configuration.

type PromptPoint 

type PromptPoint struct {
	X, Y  int
	Label PointLabel
}

PromptPoint defines a coordinate and its intent.

type Segmentation 

type Segmentation struct {
	// Mask is represented as a standard grey/binary image.Image scaled to the original image dimensions.
	// Pixels inside the mask have value 255, and outside have value 0.
	Mask     image.Image
	IoUScore float32
}

Segmentation represents a single valid segmentation mask hypothesis generated by the decoder.

type SegmentationOptions 

type SegmentationOptions struct {
	Points []PromptPoint
	Boxes  []PromptBox

	// MultiMaskOutput controls whether the model returns multiple ambiguous masks
	// (e.g., sub-part, part, whole) or just the single best mask.
	MultiMaskOutput bool
}

SegmentationOptions consolidates the inputs for the lightweight mask decoder.

type Segmenter 

type Segmenter struct {
	// contains filtered or unexported fields
}

Segmenter uses SAM2 model to segment images.

func NewSegmenter 

func NewSegmenter(backend compute.Backend, modelObj *Model) (*Segmenter, error)

NewSegmenter initializes a predictor using the underlying SAM2 model definition.

func (*Segmenter) Segment 

func (p *Segmenter) Segment(img image.Image, options *SegmentationOptions) ([]Segmentation, error)

Segment infers the segmentation masks for the given image and prompts.

type VisionConfig 

type VisionConfig struct {
	BackboneChannelList  []int   `json:"backbone_channel_list"`
	BackboneFeatureSizes [][]int `json:"backbone_feature_sizes"`
	// FPNHiddenSize is the hidden dimension of the Feature Pyramid Network (FPN) neck.
	FPNHiddenSize int `json:"fpn_hidden_size"`
	// FPNKernelSize is the kernel size used in the Feature Pyramid Network (FPN) neck.
	FPNKernelSize int `json:"fpn_kernel_size"`
	// FPNStride is the stride used in the Feature Pyramid Network (FPN) neck.
	FPNStride int `json:"fpn_stride"`
	// FPNPadding is the padding used in the Feature Pyramid Network (FPN) neck.
	FPNPadding int `json:"fpn_padding"`
	// FPNTopDownLevels specifies the levels of the Feature Pyramid Network (FPN) involved in top-down connections.
	FPNTopDownLevels []int          `json:"fpn_top_down_levels"`
	NumFeatureLevels int            `json:"num_feature_levels"`
	HiddenAct        string         `json:"hidden_act"`
	LayerNormEps     float64        `json:"layer_norm_eps"`
	BackboneConfig   HieraDetConfig `json:"backbone_config"`
}

VisionConfig represents the vision encoder configuration. It includes parameters for the Feature Pyramid Network (FPN) neck.

Source Files

View all Source files

Directories

Path Synopsis

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.