onnx

Documentation

Overview

Package onnx provides a pure-Go wrapper around the ONNX Runtime shared library for AIGC image detection. Zero CGO dependency.

Index

• Constants
• func DefaultLibPath(modelsDir string) (string, error)
• func DefaultModelPath(modelsDir string) string
• func Preprocess(img image.Image, targetSize int) []float32
• func Softmax(logits []float32) []float32
• type Detector
  • func NewDetector(libPath, modelPath string) (*Detector, error)
  • func (d *Detector) Close()
  • func (d *Detector) Detect(img image.Image) (*Result, error)
  • func (d *Detector) DetectFile(path string) (*Result, error)
  • func (d *Detector) ModelPath() string
• type Result

Constants

const (
	// ModelInputName is the expected ONNX model input tensor name.
	ModelInputName = "pixel_values"
	// ModelOutputName is the expected ONNX model output tensor name.
	ModelOutputName = "logits"
	// ModelInputSize is the expected image size (width/height) in pixels.
	ModelInputSize = 224
	// ModelChannels is the expected number of color channels.
	ModelChannels = 3
)

Functions

func DefaultLibPath

func DefaultLibPath(modelsDir string) (string, error)

DefaultLibPath returns the expected path for the ONNX Runtime shared library. Prefers GPU variant, falls back to CPU.

func DefaultModelPath

func DefaultModelPath(modelsDir string) string

DefaultModelPath returns the expected path for the ONNX model file.

func Preprocess

func Preprocess(img image.Image, targetSize int) []float32

Preprocess converts an image to a normalized float32 tensor suitable for the AIGC detection model. It resizes to targetSize x targetSize and normalizes pixel values to [0, 1].

func Softmax

func Softmax(logits []float32) []float32

Softmax applies the softmax function to the input logits. Returns probabilities that sum to 1.0.

Types

type Detector

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

Detector manages the ONNX Runtime lifecycle and inference session.

func NewDetector

func NewDetector(libPath, modelPath string) (*Detector, error)

NewDetector creates a new ONNX Detector, loading the runtime and model. Call Close() when done to release resources.

func (*Detector) Close

func (d *Detector) Close()

Close releases all ONNX Runtime resources.

func (*Detector) Detect

func (d *Detector) Detect(img image.Image) (*Result, error)

Detect runs AIGC detection on the given image. The image must be decoded (can be any format supported by image.Decode).

func (*Detector) DetectFile

func (d *Detector) DetectFile(path string) (*Result, error)

DetectFile loads an image file and runs AIGC detection.

func (*Detector) ModelPath

func (d *Detector) ModelPath() string

ModelPath returns the path to the ONNX model file used by this detector.

type Result

type Result struct {
	AIGenRate float64 // probability of being AI-generated (softmax output)
}

Result holds the AIGC detection inference result. AIGenRate is the probability (0-1) that the image is AI-generated.