first commit

pkg/matting/matting.go

@@ -0,0 +1,109 @@
+// Package matting 封装图片抠图服务调用，将衣物图片上传到 BiRefNet 抠图服务，
+// 返回透明底 PNG 图片字节。
+package matting
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"mime/multipart"
+	"net/http"
+	"time"
+
+	"git.lzcwy.cn/zhengchiliu/zt-closet-backend/global"
+)
+
+const (
+	mattingPath = "/matting"
+	healthzPath = "/healthz"
+)
+
+// Client 抠图服务客户端。
+type Client struct {
+	baseURL    string
+	httpClient *http.Client
+}
+
+// NewClient 从全局配置创建客户端。
+func NewClient() *Client {
+	return NewClientWithAddr(global.MattingModelSetting.Host, global.MattingModelSetting.Port)
+}
+
+// NewClientWithAddr 使用指定 Host/Port 创建客户端，便于测试与解耦。
+func NewClientWithAddr(host string, port int) *Client {
+	return &Client{
+		baseURL: fmt.Sprintf("http://%s:%d", host, port),
+		httpClient: &http.Client{
+			Timeout: 120 * time.Second,
+		},
+	}
+}
+
+// Matting 调用抠图服务，将图片上传并返回透明底 PNG 字节。
+func (c *Client) Matting(ctx context.Context, imageData []byte) ([]byte, error) {
+	var body bytes.Buffer
+	writer := multipart.NewWriter(&body)
+
+	part, err := writer.CreateFormFile("file", "image.png")
+	if err != nil {
+		return nil, fmt.Errorf("matting create form file: %w", err)
+	}
+	if _, err := part.Write(imageData); err != nil {
+		return nil, fmt.Errorf("matting write image: %w", err)
+	}
+	if err := writer.Close(); err != nil {
+		return nil, fmt.Errorf("matting close multipart: %w", err)
+	}
+
+	request, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+mattingPath, &body)
+	if err != nil {
+		return nil, fmt.Errorf("matting create request: %w", err)
+	}
+	request.Header.Set("Content-Type", writer.FormDataContentType())
+
+	response, err := c.httpClient.Do(request)
+	if err != nil {
+		return nil, fmt.Errorf("matting do request: %w", err)
+	}
+	defer response.Body.Close()
+
+	data, err := io.ReadAll(response.Body)
+	if err != nil {
+		return nil, fmt.Errorf("matting read response: %w", err)
+	}
+
+	if response.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("matting api error: status=%d, body=%s", response.StatusCode, summarize(data))
+	}
+
+	return data, nil
+}
+
+// HealthCheck 检查抠图服务是否就绪。
+func (c *Client) HealthCheck(ctx context.Context) error {
+	request, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+healthzPath, nil)
+	if err != nil {
+		return fmt.Errorf("matting create healthz request: %w", err)
+	}
+
+	response, err := c.httpClient.Do(request)
+	if err != nil {
+		return fmt.Errorf("matting healthz request: %w", err)
+	}
+	defer response.Body.Close()
+
+	if response.StatusCode != http.StatusOK {
+		return fmt.Errorf("matting healthz status: %d", response.StatusCode)
+	}
+	return nil
+}
+
+// summarize 截断响应体摘要，避免错误信息打印大体积二进制。
+func summarize(data []byte) string {
+	const maxLen = 512
+	if len(data) > maxLen {
+		return string(data[:maxLen]) + "..."
+	}
+	return string(data)
+}

pkg/matting/matting_test.go

@@ -0,0 +1,53 @@
+package matting
+
+import (
+	"bytes"
+	"context"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+)
+
+// TestMatting 使用仓库根目录的 image.png 调用真实抠图服务进行端到端验证。
+func TestMatting(t *testing.T) {
+	client := NewClientWithAddr("127.0.0.1", 8001)
+	ctx, cancel := context.WithTimeout(context.Background(), 180*time.Second)
+	defer cancel()
+
+	// 前置：确认抠图服务已就绪（模型加载完成后 /healthz 才会响应）
+	if err := client.HealthCheck(ctx); err != nil {
+		t.Skipf("抠图服务未就绪，跳过真实抠图测试: %v", err)
+	}
+
+	imgPath := filepath.Join("..", "..", "image.png")
+	raw, err := os.ReadFile(imgPath)
+	if err != nil {
+		t.Fatalf("读取 image.png 失败: %v", err)
+	}
+
+	out, err := client.Matting(ctx, raw)
+	if err != nil {
+		t.Fatalf("Matting 失败: %v", err)
+	}
+
+	if len(out) == 0 {
+		t.Fatal("抠图返回空结果")
+	}
+
+	pngMagic := []byte{0x89, 'P', 'N', 'G', 0x0D, 0x0A, 0x1A, 0x0A}
+	if len(out) < len(pngMagic) || !bytes.Equal(out[:len(pngMagic)], pngMagic) {
+		preview := out
+		if len(preview) > len(pngMagic) {
+			preview = preview[:len(pngMagic)]
+		}
+		t.Fatalf("返回结果不是 PNG，前 %d 字节: % x", len(preview), preview)
+	}
+
+	// 保存结果便于人工核验
+	resultPath := filepath.Join(t.TempDir(), "matting_result.png")
+	if err := os.WriteFile(resultPath, out, 0o644); err != nil {
+		t.Fatalf("保存抠图结果失败: %v", err)
+	}
+	t.Logf("抠图成功，输入 %d 字节 -> 输出 %d 字节，结果已保存到 %s", len(raw), len(out), resultPath)
+}