weijuesen
91 lines
2.0 KiB
Go
91 lines
2.0 KiB
Go
package service
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
// AIBBox 检测框
|
|
type AIBBox struct {
|
|
X float64 `json:"x"`
|
|
Y float64 `json:"y"`
|
|
W float64 `json:"w"`
|
|
H float64 `json:"h"`
|
|
}
|
|
|
|
// AIDetection 单条检测结果
|
|
type AIDetection struct {
|
|
BBox AIBBox `json:"bbox"`
|
|
ClassName string `json:"class"`
|
|
Confidence float64 `json:"confidence"`
|
|
}
|
|
|
|
// AIDetectResponse /detect 响应
|
|
type AIDetectResponse struct {
|
|
Model string `json:"model"`
|
|
Detections []AIDetection `json:"detections"`
|
|
}
|
|
|
|
// AIClient ai-service HTTP 客户端
|
|
type AIClient struct {
|
|
baseURL string
|
|
client *http.Client
|
|
}
|
|
|
|
// NewAIClient 创建 AI 客户端
|
|
func NewAIClient(baseURL string) *AIClient {
|
|
return &AIClient{
|
|
baseURL: strings.TrimSuffix(baseURL, "/"),
|
|
client: &http.Client{Timeout: 30 * time.Second},
|
|
}
|
|
}
|
|
|
|
// Detect 上传图片到 ai-service /detect，返回检测结果（无状态接口，天然幂等）
|
|
func (c *AIClient) Detect(ctx context.Context, imageBytes []byte, filename string) (*AIDetectResponse, error) {
|
|
var body bytes.Buffer
|
|
writer := multipart.NewWriter(&body)
|
|
part, err := writer.CreateFormFile("file", filename)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if _, err := part.Write(imageBytes); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := writer.Close(); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/detect", &body)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
req.Header.Set("Content-Type", writer.FormDataContentType())
|
|
|
|
resp, err := c.client.Do(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
data, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if resp.StatusCode != http.StatusOK {
|
|
return nil, fmt.Errorf("ai-service /detect 返回 %d: %s", resp.StatusCode, string(data))
|
|
}
|
|
|
|
var out AIDetectResponse
|
|
if err := json.Unmarshal(data, &out); err != nil {
|
|
return nil, err
|
|
}
|
|
return &out, nil
|
|
}
|