moteva/server/internal/infrastructure/backgroundremoval/onnxruntime.go

package backgroundremoval

import (
	"bytes"
	"context"
	"fmt"
	"image"
	"image/color"
	imagedraw "image/draw"
	"image/png"
	"math"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"

	"img_infinite_canvas/internal/domain/design"

	"github.com/zeromicro/go-zero/core/logx"
	"golang.org/x/image/draw"

	ort "github.com/yalue/onnxruntime_go"
)

const (
	defaultModelInputSize = 1024
	minAutoInputSize      = 256
	maxAutoInputSize      = 2048
	inputSizeMultiple     = 32
	defaultInputName      = "input"
	defaultOutputName     = "output"
	defaultTimeout        = 5 * time.Minute
)

var ortEnvironmentMu sync.Mutex

type ONNXRuntimeOptions struct {
	ModelPath         string
	SharedLibraryPath string
	ExecutionProvider string
	InputSize         int
	TimeoutSeconds    int
	InputName         string
	OutputName        string
}

type ONNXRuntimeRemover struct {
	modelPath         string
	sharedLibraryPath string
	executionProvider string
	inputSize         int
	timeout           time.Duration
	inputName         string
	outputName        string

	mu             sync.Mutex
	session        *ort.DynamicAdvancedSession
	providerName   string
	modelInputSize int
}

func NewONNXRuntimeRemover(opts ONNXRuntimeOptions) *ONNXRuntimeRemover {
	inputSize := opts.InputSize
	if inputSize < 0 {
		inputSize = 0
	}
	timeout := time.Duration(opts.TimeoutSeconds) * time.Second
	if timeout <= 0 {
		timeout = defaultTimeout
	}
	inputName := strings.TrimSpace(opts.InputName)
	if inputName == "" {
		inputName = defaultInputName
	}
	outputName := strings.TrimSpace(opts.OutputName)
	if outputName == "" {
		outputName = defaultOutputName
	}
	return &ONNXRuntimeRemover{
		modelPath:         resolvePath(opts.ModelPath, "model/rmbg1.4.onnx"),
		sharedLibraryPath: resolveSharedLibraryPath(opts.SharedLibraryPath),
		executionProvider: normalizeExecutionProvider(opts.ExecutionProvider),
		inputSize:         inputSize,
		timeout:           timeout,
		inputName:         inputName,
		outputName:        outputName,
	}
}

func (r *ONNXRuntimeRemover) RemoveBackground(ctx context.Context, req design.BackgroundRemovalRequest) (design.BackgroundRemoval, error) {
	if len(req.Image) == 0 {
		return design.BackgroundRemoval{}, fmt.Errorf("%w: image is required", design.ErrInvalidInput)
	}
	if r.timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, r.timeout)
		defer cancel()
	}
	if err := ctx.Err(); err != nil {
		return design.BackgroundRemoval{}, err
	}

	source, _, err := image.Decode(bytes.NewReader(req.Image))
	if err != nil {
		return design.BackgroundRemoval{}, err
	}
	sourceBounds := source.Bounds()
	r.mu.Lock()
	if err := r.ensureSession(); err != nil {
		r.mu.Unlock()
		return design.BackgroundRemoval{}, err
	}
	inputSize := r.effectiveInputSize(req, sourceBounds.Dx(), sourceBounds.Dy())
	r.mu.Unlock()

	inputData, sourceRGBA := preprocessImage(source, inputSize)
	inputTensor, err := ort.NewTensor(ort.NewShape(1, 3, int64(inputSize), int64(inputSize)), inputData)
	if err != nil {
		return design.BackgroundRemoval{}, err
	}
	defer inputTensor.Destroy()

	outputs := []ort.Value{nil}
	r.mu.Lock()
	err = r.session.Run([]ort.Value{inputTensor}, outputs)
	r.mu.Unlock()
	if err != nil {
		return design.BackgroundRemoval{}, err
	}
	if outputs[0] != nil {
		defer outputs[0].Destroy()
	}
	if err := ctx.Err(); err != nil {
		return design.BackgroundRemoval{}, err
	}

	outputTensor, ok := outputs[0].(*ort.Tensor[float32])
	if !ok {
		return design.BackgroundRemoval{}, fmt.Errorf("rmbg output tensor has unsupported type %T", outputs[0])
	}
	mask := postprocessMask(outputTensor.GetData(), outputTensor.GetShape(), sourceRGBA.Bounds().Dx(), sourceRGBA.Bounds().Dy())
	result := applyAlphaMask(sourceRGBA, mask)

	var output bytes.Buffer
	if err := png.Encode(&output, result); err != nil {
		return design.BackgroundRemoval{}, err
	}
	return design.BackgroundRemoval{
		Image:       output.Bytes(),
		ContentType: "image/png",
		Width:       result.Bounds().Dx(),
		Height:      result.Bounds().Dy(),
	}, nil
}

func (r *ONNXRuntimeRemover) ensureSession() error {
	if r.session != nil {
		return nil
	}
	if err := ensureORTEnvironment(r.sharedLibraryPath); err != nil {
		return err
	}
	if inputSize, err := fixedModelInputSize(r.modelPath, r.inputName); err == nil && inputSize > 0 {
		r.modelInputSize = inputSize
	} else if err != nil {
		logx.Errorf("read background removal model input size failed: %v", err)
	}
	session, provider, err := createSessionWithFallback(r.modelPath, r.inputName, r.outputName, r.executionProvider)
	if err != nil {
		return err
	}
	r.session = session
	r.providerName = provider
	logx.Infof("background removal ONNX Runtime provider: %s", provider)
	return nil
}

func (r *ONNXRuntimeRemover) effectiveInputSize(req design.BackgroundRemovalRequest, imageWidth int, imageHeight int) int {
	if r.modelInputSize > 0 {
		return r.modelInputSize
	}
	if req.InputSize > 0 {
		return normalizeInputSize(req.InputSize)
	}
	if r.inputSize > 0 {
		return normalizeInputSize(r.inputSize)
	}
	if req.Width > 0 || req.Height > 0 {
		return inputSizeFromImageSize(req.Width, req.Height)
	}
	return inputSizeFromImageSize(imageWidth, imageHeight)
}

func fixedModelInputSize(modelPath string, inputName string) (int, error) {
	inputs, _, err := ort.GetInputOutputInfo(modelPath)
	if err != nil {
		return 0, err
	}
	if len(inputs) == 0 {
		return 0, nil
	}
	selected := inputs[0]
	for _, input := range inputs {
		if input.Name == inputName {
			selected = input
			break
		}
	}
	shape := selected.Dimensions
	if len(shape) < 4 {
		return 0, nil
	}
	height := int(shape[len(shape)-2])
	width := int(shape[len(shape)-1])
	if height <= 0 || width <= 0 {
		return 0, nil
	}
	if height > width {
		return height, nil
	}
	return width, nil
}

func ensureORTEnvironment(sharedLibraryPath string) error {
	ortEnvironmentMu.Lock()
	defer ortEnvironmentMu.Unlock()
	if ort.IsInitialized() {
		return nil
	}
	if sharedLibraryPath != "" {
		ort.SetSharedLibraryPath(sharedLibraryPath)
	}
	if err := ort.InitializeEnvironment(); err != nil {
		return fmt.Errorf("initialize ONNX Runtime: %w", err)
	}
	return nil
}

func createSessionWithFallback(modelPath string, inputName string, outputName string, provider string) (*ort.DynamicAdvancedSession, string, error) {
	attempts := executionProviderAttempts(provider)
	var failures []string
	for _, attempt := range attempts {
		session, err := createSession(modelPath, inputName, outputName, attempt)
		if err == nil {
			return session, attempt, nil
		}
		failures = append(failures, fmt.Sprintf("%s: %v", attempt, err))
		if provider != "auto" && provider != "gpu" {
			break
		}
	}
	return nil, "", fmt.Errorf("create ONNX Runtime session failed (%s)", strings.Join(failures, "; "))
}

func createSession(modelPath string, inputName string, outputName string, provider string) (*ort.DynamicAdvancedSession, error) {
	options, err := ort.NewSessionOptions()
	if err != nil {
		return nil, err
	}
	defer options.Destroy()
	if err := options.SetIntraOpNumThreads(0); err != nil {
		return nil, err
	}
	if err := appendExecutionProvider(options, provider); err != nil {
		return nil, err
	}
	return ort.NewDynamicAdvancedSession(modelPath, []string{inputName}, []string{outputName}, options)
}

func appendExecutionProvider(options *ort.SessionOptions, provider string) error {
	switch provider {
	case "cpu":
		return nil
	case "cuda":
		cudaOptions, err := ort.NewCUDAProviderOptions()
		if err != nil {
			return err
		}
		defer cudaOptions.Destroy()
		if err := cudaOptions.Update(map[string]string{"device_id": "0"}); err != nil {
			return err
		}
		return options.AppendExecutionProviderCUDA(cudaOptions)
	case "coreml":
		return options.AppendExecutionProviderCoreMLV2(map[string]string{
			"MLComputeUnits": "ALL",
		})
	case "directml":
		return options.AppendExecutionProviderDirectML(0)
	default:
		return fmt.Errorf("%w: unsupported ONNX execution provider %q", design.ErrInvalidInput, provider)
	}
}

func executionProviderAttempts(provider string) []string {
	switch provider {
	case "cpu":
		return []string{"cpu"}
	case "cuda":
		return []string{"cuda", "cpu"}
	case "coreml":
		return []string{"coreml", "cpu"}
	case "directml":
		return []string{"directml", "cpu"}
	case "gpu", "auto":
		switch runtime.GOOS {
		case "darwin", "ios":
			return []string{"coreml", "cpu"}
		case "windows":
			return []string{"cuda", "directml", "cpu"}
		default:
			return []string{"cuda", "cpu"}
		}
	default:
		return []string{"cpu"}
	}
}

func normalizeExecutionProvider(provider string) string {
	provider = strings.ToLower(strings.TrimSpace(provider))
	provider = strings.ReplaceAll(provider, "_", "")
	provider = strings.ReplaceAll(provider, "-", "")
	switch provider {
	case "", "auto":
		return "auto"
	case "gpu":
		return "gpu"
	case "cpu", "cpuexecutionprovider":
		return "cpu"
	case "cuda", "cudaexecutionprovider":
		return "cuda"
	case "coreml", "coremlexecutionprovider":
		return "coreml"
	case "directml", "dml", "directmlexecutionprovider":
		return "directml"
	default:
		return provider
	}
}

func preprocessImage(source image.Image, inputSize int) ([]float32, *image.NRGBA) {
	sourceBounds := source.Bounds()
	sourceRGBA := image.NewNRGBA(image.Rect(0, 0, sourceBounds.Dx(), sourceBounds.Dy()))
	imagedraw.Draw(sourceRGBA, sourceRGBA.Bounds(), source, sourceBounds.Min, imagedraw.Src)

	resized := image.NewNRGBA(image.Rect(0, 0, inputSize, inputSize))
	draw.CatmullRom.Scale(resized, resized.Bounds(), sourceRGBA, sourceRGBA.Bounds(), draw.Over, nil)

	planeSize := inputSize * inputSize
	data := make([]float32, 3*planeSize)
	for y := 0; y < inputSize; y++ {
		for x := 0; x < inputSize; x++ {
			offset := resized.PixOffset(x, y)
			index := y*inputSize + x
			data[index] = float32(resized.Pix[offset])/255 - 0.5
			data[planeSize+index] = float32(resized.Pix[offset+1])/255 - 0.5
			data[2*planeSize+index] = float32(resized.Pix[offset+2])/255 - 0.5
		}
	}
	return data, sourceRGBA
}

func postprocessMask(data []float32, shape ort.Shape, width int, height int) *image.Gray {
	if len(data) == 0 || width <= 0 || height <= 0 {
		return image.NewGray(image.Rect(0, 0, maxInt(width, 1), maxInt(height, 1)))
	}
	maskWidth, maskHeight := maskDimensions(shape, len(data))
	mask := image.NewGray(image.Rect(0, 0, maskWidth, maskHeight))
	minValue, maxValue := data[0], data[0]
	for _, value := range data {
		if value < minValue {
			minValue = value
		}
		if value > maxValue {
			maxValue = value
		}
	}
	denominator := maxValue - minValue
	for y := 0; y < maskHeight; y++ {
		for x := 0; x < maskWidth; x++ {
			index := y*maskWidth + x
			value := float32(0)
			if index < len(data) && denominator > 1e-6 {
				value = (data[index] - minValue) / denominator
			}
			value = float32(math.Max(0, math.Min(1, float64(value))))
			mask.SetGray(x, y, color.Gray{Y: uint8(value*255 + 0.5)})
		}
	}
	if maskWidth == width && maskHeight == height {
		return mask
	}
	resized := image.NewGray(image.Rect(0, 0, width, height))
	draw.CatmullRom.Scale(resized, resized.Bounds(), mask, mask.Bounds(), draw.Over, nil)
	return resized
}

func maskDimensions(shape ort.Shape, dataLength int) (int, int) {
	if len(shape) >= 2 {
		width := int(shape[len(shape)-1])
		height := int(shape[len(shape)-2])
		if width > 0 && height > 0 && width*height <= dataLength {
			return width, height
		}
	}
	side := int(math.Sqrt(float64(dataLength)))
	if side > 0 && side*side == dataLength {
		return side, side
	}
	return dataLength, 1
}

func inputSizeFromImageSize(width int, height int) int {
	maxSide := width
	if height > maxSide {
		maxSide = height
	}
	if maxSide <= 0 {
		return defaultModelInputSize
	}
	return normalizeInputSize(maxSide)
}

func normalizeInputSize(size int) int {
	if size <= 0 {
		return defaultModelInputSize
	}
	if size < minAutoInputSize {
		size = minAutoInputSize
	}
	if size > maxAutoInputSize {
		size = maxAutoInputSize
	}
	if remainder := size % inputSizeMultiple; remainder != 0 {
		size += inputSizeMultiple - remainder
	}
	if size > maxAutoInputSize {
		size = maxAutoInputSize
	}
	return size
}

func applyAlphaMask(source *image.NRGBA, mask *image.Gray) *image.NRGBA {
	bounds := source.Bounds()
	result := image.NewNRGBA(bounds)
	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			sourceOffset := source.PixOffset(x, y)
			maskAlpha := mask.GrayAt(x-bounds.Min.X, y-bounds.Min.Y).Y
			sourceAlpha := source.Pix[sourceOffset+3]
			alpha := uint8((uint16(maskAlpha)*uint16(sourceAlpha) + 127) / 255)
			resultOffset := result.PixOffset(x, y)
			result.Pix[resultOffset] = source.Pix[sourceOffset]
			result.Pix[resultOffset+1] = source.Pix[sourceOffset+1]
			result.Pix[resultOffset+2] = source.Pix[sourceOffset+2]
			result.Pix[resultOffset+3] = alpha
		}
	}
	return result
}

func resolveSharedLibraryPath(configured string) string {
	for _, candidate := range []string{
		configured,
		os.Getenv("ONNXRUNTIME_SHARED_LIBRARY_PATH"),
		os.Getenv("ORT_SHARED_LIBRARY_PATH"),
		os.Getenv("ORT_DYLIB_PATH"),
		"model/libonnxruntime.dylib",
		"model/libonnxruntime.so",
		"model/onnxruntime.dll",
		"server/model/libonnxruntime.dylib",
		"server/model/libonnxruntime.so",
		"server/model/onnxruntime.dll",
		"/opt/homebrew/lib/libonnxruntime.dylib",
		"/usr/local/lib/libonnxruntime.dylib",
		"/usr/local/lib/libonnxruntime.so",
	} {
		path := resolveExistingPath(candidate)
		if path != "" {
			return path
		}
	}
	return strings.TrimSpace(configured)
}

func resolvePath(configured string, fallback string) string {
	path := resolveExistingPath(configured)
	if path != "" {
		return path
	}
	path = resolveExistingPath(fallback)
	if path != "" {
		return path
	}
	return strings.TrimSpace(firstNonEmpty(configured, fallback))
}

func resolveExistingPath(value string) string {
	value = strings.TrimSpace(value)
	if value == "" {
		return ""
	}
	candidates := []string{value}
	if !filepath.IsAbs(value) {
		for _, base := range ancestorDirs() {
			candidates = append(candidates, filepath.Join(base, value), filepath.Join(base, "server", value))
		}
	}
	for _, candidate := range candidates {
		if _, err := os.Stat(candidate); err == nil {
			if absolute, err := filepath.Abs(candidate); err == nil {
				return absolute
			}
			return candidate
		}
	}
	return ""
}

func ancestorDirs() []string {
	cwd, err := os.Getwd()
	if err != nil {
		return []string{"."}
	}
	var dirs []string
	for {
		dirs = append(dirs, cwd)
		parent := filepath.Dir(cwd)
		if parent == cwd {
			break
		}
		cwd = parent
		if len(dirs) >= 8 {
			break
		}
	}
	return dirs
}

func firstNonEmpty(values ...string) string {
	for _, value := range values {
		if strings.TrimSpace(value) != "" {
			return value
		}
	}
	return ""
}

func maxInt(a int, b int) int {
	if a > b {
		return a
	}
	return b
}