first commit
This commit is contained in:
@@ -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)
|
||||
}
|
||||
Reference in New Issue
Block a user