59 lines
1.5 KiB
Go
59 lines
1.5 KiB
Go
|
package backgroundremoval
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"context"
|
||
|
|
"image"
|
||
|
|
"image/color"
|
||
|
|
"image/png"
|
||
|
|
"os"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"img_infinite_canvas/internal/domain/design"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestONNXRuntimeRemoverIntegration(t *testing.T) {
|
||
|
|
if os.Getenv("RUN_RMBG_ONNX_INTEGRATION") != "1" {
|
||
|
|
t.Skip("set RUN_RMBG_ONNX_INTEGRATION=1 to run the local rmbg model")
|
||
|
|
}
|
||
|
|
var input bytes.Buffer
|
||
|
|
img := image.NewNRGBA(image.Rect(0, 0, 32, 32))
|
||
|
|
for y := 0; y < 32; y++ {
|
||
|
|
for x := 0; x < 32; x++ {
|
||
|
|
c := color.NRGBA{R: 240, G: 240, B: 240, A: 255}
|
||
|
|
if x >= 8 && x < 24 && y >= 8 && y < 24 {
|
||
|
|
c = color.NRGBA{R: 220, G: 40, B: 40, A: 255}
|
||
|
|
}
|
||
|
|
img.SetNRGBA(x, y, c)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if err := png.Encode(&input, img); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
|
||
|
|
remover := NewONNXRuntimeRemover(ONNXRuntimeOptions{
|
||
|
|
ModelPath: "model/rmbg1.4.onnx",
|
||
|
|
SharedLibraryPath: os.Getenv("ONNXRUNTIME_SHARED_LIBRARY_PATH"),
|
||
|
|
ExecutionProvider: "auto",
|
||
|
|
InputSize: 1024,
|
||
|
|
TimeoutSeconds: 300,
|
||
|
|
})
|
||
|
|
result, err := remover.RemoveBackground(context.Background(), design.BackgroundRemovalRequest{
|
||
|
|
Image: input.Bytes(),
|
||
|
|
ContentType: "image/png",
|
||
|
|
})
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if result.ContentType != "image/png" || len(result.Image) == 0 {
|
||
|
|
t.Fatalf("unexpected result metadata: %#v", result)
|
||
|
|
}
|
||
|
|
decoded, _, err := image.Decode(bytes.NewReader(result.Image))
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if decoded.Bounds().Dx() != 32 || decoded.Bounds().Dy() != 32 {
|
||
|
|
t.Fatalf("unexpected output size %s", decoded.Bounds())
|
||
|
|
}
|
||
|
|
}
|