289 lines
7.8 KiB
Go
289 lines
7.8 KiB
Go
|
package objectrecognition
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"context"
|
||
|
|
"encoding/base64"
|
||
|
|
"encoding/json"
|
||
|
|
"fmt"
|
||
|
|
"io"
|
||
|
|
"net/http"
|
||
|
|
"strings"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"img_infinite_canvas/internal/domain/design"
|
||
|
|
)
|
||
|
|
|
||
|
|
type ClientOptions struct {
|
||
|
|
BaseURL string
|
||
|
|
APIKey string
|
||
|
|
Model string
|
||
|
|
TimeoutSeconds int
|
||
|
|
}
|
||
|
|
|
||
|
|
type Client struct {
|
||
|
|
baseURL string
|
||
|
|
apiKey string
|
||
|
|
model string
|
||
|
|
client *http.Client
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewClient(opts ClientOptions) *Client {
|
||
|
|
if strings.TrimSpace(opts.BaseURL) == "" || strings.TrimSpace(opts.APIKey) == "" {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
timeout := time.Duration(opts.TimeoutSeconds) * time.Second
|
||
|
|
if timeout <= 0 {
|
||
|
|
timeout = 60 * time.Second
|
||
|
|
}
|
||
|
|
model := strings.TrimSpace(opts.Model)
|
||
|
|
if model == "" {
|
||
|
|
model = "gpt-5.4-mini"
|
||
|
|
}
|
||
|
|
return &Client{
|
||
|
|
baseURL: strings.TrimRight(opts.BaseURL, "/"),
|
||
|
|
apiKey: opts.APIKey,
|
||
|
|
model: model,
|
||
|
|
client: &http.Client{Timeout: timeout},
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *Client) RecognizeObject(ctx context.Context, req design.ObjectRecognitionRequest) (design.ObjectRecognition, error) {
|
||
|
|
imageURL := c.imagePayloadURL(ctx, req)
|
||
|
|
if strings.TrimSpace(imageURL) == "" {
|
||
|
|
return design.ObjectRecognition{}, fmt.Errorf("%w: image_base64 or image_url is required", design.ErrInvalidInput)
|
||
|
|
}
|
||
|
|
box := clampBox(req.BoundingBox)
|
||
|
|
lang := strings.TrimSpace(req.Lang)
|
||
|
|
if lang == "" {
|
||
|
|
lang = "zh"
|
||
|
|
}
|
||
|
|
|
||
|
|
body, err := json.Marshal(map[string]any{
|
||
|
|
"model": c.model,
|
||
|
|
"messages": []map[string]any{
|
||
|
|
{
|
||
|
|
"role": "system",
|
||
|
|
"content": strings.Join([]string{
|
||
|
|
"You are a concise object recognition engine for an image editor.",
|
||
|
|
"Identify only the visible object inside the provided normalized bounding box.",
|
||
|
|
"Return strict JSON only: {\"id\":\"\",\"label\":\"short object label\",\"confidence\":0.0,\"bbox\":{\"x\":0,\"y\":0,\"width\":0,\"height\":0}}.",
|
||
|
|
"Use the requested language for label: Chinese for zh, English for en.",
|
||
|
|
"Keep labels short and concrete. Do not describe the whole scene.",
|
||
|
|
"Return the input bbox unless you can confidently tighten it to the same object.",
|
||
|
|
}, " "),
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"role": "user",
|
||
|
|
"content": []map[string]any{
|
||
|
|
{
|
||
|
|
"type": "text",
|
||
|
|
"text": fmt.Sprintf("Recognize the object inside bbox x=%.6f y=%.6f width=%.6f height=%.6f. Language: %s. Extra hint: %s", box.X, box.Y, box.Width, box.Height, lang, strings.TrimSpace(req.Prompt)),
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"type": "image_url",
|
||
|
|
"image_url": map[string]any{
|
||
|
|
"url": imageURL,
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
"response_format": map[string]any{"type": "json_object"},
|
||
|
|
"max_tokens": 220,
|
||
|
|
"temperature": 0.1,
|
||
|
|
})
|
||
|
|
if err != nil {
|
||
|
|
return design.ObjectRecognition{}, err
|
||
|
|
}
|
||
|
|
|
||
|
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/chat/completions", bytes.NewReader(body))
|
||
|
|
if err != nil {
|
||
|
|
return design.ObjectRecognition{}, err
|
||
|
|
}
|
||
|
|
httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
|
||
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
||
|
|
|
||
|
|
resp, err := c.client.Do(httpReq)
|
||
|
|
if err != nil {
|
||
|
|
return design.ObjectRecognition{}, err
|
||
|
|
}
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
respBody, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
|
||
|
|
if err != nil {
|
||
|
|
return design.ObjectRecognition{}, err
|
||
|
|
}
|
||
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||
|
|
return design.ObjectRecognition{}, fmt.Errorf("object recognition failed: %s: %s", resp.Status, trimErrorBody(respBody))
|
||
|
|
}
|
||
|
|
|
||
|
|
var payload struct {
|
||
|
|
Choices []struct {
|
||
|
|
Message struct {
|
||
|
|
Content json.RawMessage `json:"content"`
|
||
|
|
} `json:"message"`
|
||
|
|
} `json:"choices"`
|
||
|
|
}
|
||
|
|
if err := json.Unmarshal(respBody, &payload); err != nil {
|
||
|
|
return design.ObjectRecognition{}, err
|
||
|
|
}
|
||
|
|
if len(payload.Choices) == 0 {
|
||
|
|
return design.ObjectRecognition{}, fmt.Errorf("object recognition returned no choices")
|
||
|
|
}
|
||
|
|
result, err := parseObjectRecognition(decodeChatContent(payload.Choices[0].Message.Content), box, lang)
|
||
|
|
if err != nil {
|
||
|
|
return design.ObjectRecognition{}, err
|
||
|
|
}
|
||
|
|
return result, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *Client) imagePayloadURL(ctx context.Context, req design.ObjectRecognitionRequest) string {
|
||
|
|
imageBase64 := strings.TrimSpace(req.ImageBase64)
|
||
|
|
if imageBase64 != "" {
|
||
|
|
if strings.HasPrefix(imageBase64, "data:image/") {
|
||
|
|
return imageBase64
|
||
|
|
}
|
||
|
|
return "data:image/jpeg;base64," + imageBase64
|
||
|
|
}
|
||
|
|
|
||
|
|
imageURL := strings.TrimSpace(req.ImageURL)
|
||
|
|
if imageURL == "" || strings.HasPrefix(imageURL, "data:image/") {
|
||
|
|
return imageURL
|
||
|
|
}
|
||
|
|
if !strings.HasPrefix(imageURL, "http://") && !strings.HasPrefix(imageURL, "https://") {
|
||
|
|
return imageURL
|
||
|
|
}
|
||
|
|
|
||
|
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, imageURL, nil)
|
||
|
|
if err != nil {
|
||
|
|
return imageURL
|
||
|
|
}
|
||
|
|
resp, err := c.client.Do(httpReq)
|
||
|
|
if err != nil {
|
||
|
|
return imageURL
|
||
|
|
}
|
||
|
|
defer resp.Body.Close()
|
||
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||
|
|
return imageURL
|
||
|
|
}
|
||
|
|
|
||
|
|
data, err := io.ReadAll(io.LimitReader(resp.Body, 20<<20))
|
||
|
|
if err != nil || len(data) == 0 {
|
||
|
|
return imageURL
|
||
|
|
}
|
||
|
|
contentType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0])
|
||
|
|
if contentType == "" || !strings.HasPrefix(contentType, "image/") {
|
||
|
|
contentType = http.DetectContentType(data)
|
||
|
|
}
|
||
|
|
if !strings.HasPrefix(contentType, "image/") {
|
||
|
|
return imageURL
|
||
|
|
}
|
||
|
|
return "data:" + contentType + ";base64," + base64.StdEncoding.EncodeToString(data)
|
||
|
|
}
|
||
|
|
|
||
|
|
func parseObjectRecognition(content string, fallbackBox design.NormalizedBox, lang string) (design.ObjectRecognition, error) {
|
||
|
|
content = strings.TrimSpace(content)
|
||
|
|
if content == "" {
|
||
|
|
return design.ObjectRecognition{}, fmt.Errorf("object recognition returned empty content")
|
||
|
|
}
|
||
|
|
if strings.HasPrefix(content, "```") {
|
||
|
|
content = strings.Trim(content, "`")
|
||
|
|
content = strings.TrimSpace(strings.TrimPrefix(content, "json"))
|
||
|
|
}
|
||
|
|
if start := strings.Index(content, "{"); start >= 0 {
|
||
|
|
if end := strings.LastIndex(content, "}"); end >= start {
|
||
|
|
content = content[start : end+1]
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
var parsed struct {
|
||
|
|
ID string `json:"id"`
|
||
|
|
Label string `json:"label"`
|
||
|
|
Confidence float64 `json:"confidence"`
|
||
|
|
BBox design.NormalizedBox `json:"bbox"`
|
||
|
|
Box design.NormalizedBox `json:"box"`
|
||
|
|
}
|
||
|
|
if err := json.Unmarshal([]byte(content), &parsed); err != nil {
|
||
|
|
return design.ObjectRecognition{}, err
|
||
|
|
}
|
||
|
|
label := strings.TrimSpace(parsed.Label)
|
||
|
|
if label == "" {
|
||
|
|
if strings.HasPrefix(strings.ToLower(strings.TrimSpace(lang)), "en") {
|
||
|
|
label = "selected object"
|
||
|
|
} else {
|
||
|
|
label = "选中对象"
|
||
|
|
}
|
||
|
|
}
|
||
|
|
box := parsed.BBox
|
||
|
|
if box.Width <= 0 || box.Height <= 0 {
|
||
|
|
box = parsed.Box
|
||
|
|
}
|
||
|
|
if box.Width <= 0 || box.Height <= 0 {
|
||
|
|
box = fallbackBox
|
||
|
|
}
|
||
|
|
return design.ObjectRecognition{
|
||
|
|
ID: strings.TrimSpace(parsed.ID),
|
||
|
|
Label: label,
|
||
|
|
Confidence: clampFloat(parsed.Confidence, 0, 1),
|
||
|
|
BBox: clampBox(box),
|
||
|
|
}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func decodeChatContent(raw json.RawMessage) string {
|
||
|
|
var text string
|
||
|
|
if err := json.Unmarshal(raw, &text); err == nil {
|
||
|
|
return text
|
||
|
|
}
|
||
|
|
var parts []struct {
|
||
|
|
Type string `json:"type"`
|
||
|
|
Text string `json:"text"`
|
||
|
|
}
|
||
|
|
if err := json.Unmarshal(raw, &parts); err == nil {
|
||
|
|
var builder strings.Builder
|
||
|
|
for _, part := range parts {
|
||
|
|
if strings.TrimSpace(part.Text) == "" {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
if builder.Len() > 0 {
|
||
|
|
builder.WriteString("\n")
|
||
|
|
}
|
||
|
|
builder.WriteString(part.Text)
|
||
|
|
}
|
||
|
|
return builder.String()
|
||
|
|
}
|
||
|
|
return string(raw)
|
||
|
|
}
|
||
|
|
|
||
|
|
func trimErrorBody(body []byte) string {
|
||
|
|
value := strings.TrimSpace(string(body))
|
||
|
|
if value == "" {
|
||
|
|
return "empty error body"
|
||
|
|
}
|
||
|
|
if len(value) > 1200 {
|
||
|
|
return value[:1200]
|
||
|
|
}
|
||
|
|
return value
|
||
|
|
}
|
||
|
|
|
||
|
|
func clampBox(box design.NormalizedBox) design.NormalizedBox {
|
||
|
|
width := clampFloat(box.Width, 0, 1)
|
||
|
|
height := clampFloat(box.Height, 0, 1)
|
||
|
|
return design.NormalizedBox{
|
||
|
|
X: clampFloat(box.X, 0, 1-width),
|
||
|
|
Y: clampFloat(box.Y, 0, 1-height),
|
||
|
|
Width: width,
|
||
|
|
Height: height,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func clampFloat(value float64, minValue float64, maxValue float64) float64 {
|
||
|
|
if value < minValue {
|
||
|
|
return minValue
|
||
|
|
}
|
||
|
|
if value > maxValue {
|
||
|
|
return maxValue
|
||
|
|
}
|
||
|
|
return value
|
||
|
|
}
|