85 lines
1.9 KiB
Go
85 lines
1.9 KiB
Go
package service
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"fmt"
|
||
"io"
|
||
"mime/multipart"
|
||
"net/http"
|
||
"strings"
|
||
"time"
|
||
)
|
||
|
||
// SegmentHTTPClient 自部署抠图 HTTP 服务（如 rembg s → POST /api/remove）
|
||
type SegmentHTTPClient struct {
|
||
endpoint string
|
||
client *http.Client
|
||
}
|
||
|
||
func NewSegmentHTTPClient(endpoint string) *SegmentHTTPClient {
|
||
return &SegmentHTTPClient{
|
||
endpoint: strings.TrimSpace(endpoint),
|
||
client: &http.Client{Timeout: 120 * time.Second},
|
||
}
|
||
}
|
||
|
||
func (c *SegmentHTTPClient) enabled() bool {
|
||
return c.endpoint != ""
|
||
}
|
||
|
||
// RemoveBackground 上传图片字节，返回透明 PNG
|
||
func (c *SegmentHTTPClient) RemoveBackground(ctx context.Context, imageData []byte, contentType string) ([]byte, error) {
|
||
if !c.enabled() {
|
||
return nil, fmt.Errorf("抠图 HTTP 地址未配置")
|
||
}
|
||
if len(imageData) == 0 {
|
||
return nil, fmt.Errorf("图片为空")
|
||
}
|
||
|
||
ext := "jpg"
|
||
if strings.Contains(strings.ToLower(contentType), "png") {
|
||
ext = "png"
|
||
}
|
||
|
||
var buf bytes.Buffer
|
||
w := multipart.NewWriter(&buf)
|
||
part, err := w.CreateFormFile("image", "upload."+ext)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if _, err := part.Write(imageData); err != nil {
|
||
return nil, err
|
||
}
|
||
_ = w.WriteField("scene", "portrait")
|
||
if err := w.Close(); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.endpoint, &buf)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
req.Header.Set("Content-Type", w.FormDataContentType())
|
||
|
||
res, err := c.client.Do(req)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer res.Body.Close()
|
||
|
||
raw, err := io.ReadAll(io.LimitReader(res.Body, maxSegmentImageBytes*4))
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if res.StatusCode < 200 || res.StatusCode >= 300 {
|
||
return nil, fmt.Errorf("抠图 HTTP %d: %s", res.StatusCode, truncate(string(raw), 240))
|
||
}
|
||
|
||
ct := strings.ToLower(res.Header.Get("Content-Type"))
|
||
if strings.Contains(ct, "image/") {
|
||
return raw, nil
|
||
}
|
||
return nil, fmt.Errorf("抠图 HTTP 返回非图片内容")
|
||
}
|