weijuesen
126 lines
3.0 KiB
Go
126 lines
3.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"`
|
|
ModelVersion string `json:"modelVersion"`
|
|
IsMock bool `json:"isMock"`
|
|
Status string `json:"status"`
|
|
AbnormalProbability float64 `json:"abnormalProbability"`
|
|
Detections []AIDetection `json:"detections"`
|
|
}
|
|
|
|
// AIDetectionStatus 从检测结果归纳 AI 状态；空检测或 unknown 不当作 healthy。
|
|
func AIDetectionStatus(detections []AIDetection) string {
|
|
if len(detections) == 0 {
|
|
return "unknown"
|
|
}
|
|
healthySeen := false
|
|
unknownSeen := false
|
|
for _, d := range detections {
|
|
class := strings.ToLower(strings.TrimSpace(d.ClassName))
|
|
switch class {
|
|
case "healthy":
|
|
healthySeen = true
|
|
case "", "unknown":
|
|
unknownSeen = true
|
|
default:
|
|
return "abnormal"
|
|
}
|
|
}
|
|
if healthySeen && !unknownSeen {
|
|
return "healthy"
|
|
}
|
|
return "unknown"
|
|
}
|
|
|
|
// 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) {
|
|
start := time.Now()
|
|
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)
|
|
latency := time.Since(start)
|
|
if err != nil {
|
|
RecordDependency(DependencyAI, err, latency)
|
|
return nil, err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
data, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if resp.StatusCode != http.StatusOK {
|
|
err := fmt.Errorf("ai-service /detect 返回 %d: %s", resp.StatusCode, string(data))
|
|
RecordDependency(DependencyAI, err, latency)
|
|
return nil, err
|
|
}
|
|
|
|
var out AIDetectResponse
|
|
if err := json.Unmarshal(data, &out); err != nil {
|
|
RecordDependency(DependencyAI, err, latency)
|
|
return nil, err
|
|
}
|
|
RecordDependency(DependencyAI, nil, latency)
|
|
return &out, nil
|
|
}
|