239 lines
8.1 KiB
Go
239 lines
8.1 KiB
Go
package common
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/base64"
|
||
"encoding/json"
|
||
"fmt"
|
||
"image"
|
||
"image/color"
|
||
"image/jpeg"
|
||
"image/png"
|
||
"io"
|
||
"net/http"
|
||
"strings"
|
||
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
)
|
||
|
||
// LocalAi RF-DETR 检测服务客户端（local-ai 兼容 /v1/detection 协议）：
|
||
// POST {baseUrl}/v1/detection，body {"model","image":"data:image/jpeg;base64,...","threshold"}，
|
||
// 响应 {"detections":[{"x","y","width","height","confidence","class_name"}]}，
|
||
// 坐标单位 = 提交图片像素，且 x/y 为框中心坐标（RF-DETR 输出约定，2026-09-02 实验确认）。
|
||
type LocalAi struct {
|
||
BaseUrl string
|
||
Model string
|
||
Threshold float64
|
||
ConfConfirmed float64
|
||
// 重叠去重阈值（minIoU = 交叠面积/两框较小面积）：与高置信框重叠超过该值的框剔除（同目标只留一个）。
|
||
// 用 minIoU 而非 IoU：RF-DETR 对同一目标常输出一大一小两个框，标准 IoU 可能仅 0.3~0.5 而漏杀，
|
||
// 大框套小框时小框被覆盖比例高，minIoU 能命中；相邻目标两框互有外露，minIoU 通常 < 0.3。
|
||
OverlapThreshold float64
|
||
// 提交前整图等比缩放到的最长边（RF-DETR 对小图更稳）
|
||
InputSize int
|
||
}
|
||
|
||
// Detection RF-DETR 单目标检测结果；Detect 返回前已映射为原图像素左上角（X/Y）+ 宽高
|
||
type Detection struct {
|
||
X float64 `json:"x"`
|
||
Y float64 `json:"y"`
|
||
Width float64 `json:"width"`
|
||
Height float64 `json:"height"`
|
||
Confidence float64 `json:"confidence"`
|
||
ClassName string `json:"class_name"`
|
||
}
|
||
|
||
// LocalAiClient 当前配置的标注服务客户端（未配置 baseUrl 返回 nil，调用方判 CodeLocalAiNotConfigured）
|
||
func LocalAiClient(ctx context.Context) *LocalAi {
|
||
base := g.Cfg().MustGet(ctx, "localAi.baseUrl").String()
|
||
if base == "" {
|
||
return nil
|
||
}
|
||
return &LocalAi{
|
||
BaseUrl: strings.TrimRight(base, "/"),
|
||
Model: g.Cfg().MustGet(ctx, "localAi.model", "rfdetr-xlarge").String(),
|
||
Threshold: g.Cfg().MustGet(ctx, "localAi.threshold", 0.3).Float64(),
|
||
ConfConfirmed: g.Cfg().MustGet(ctx, "localAi.confConfirmed", 0.5).Float64(),
|
||
OverlapThreshold: g.Cfg().MustGet(ctx, "localAi.overlapThreshold", 0.3).Float64(),
|
||
InputSize: g.Cfg().MustGet(ctx, "localAi.inputSize", 700).Int(),
|
||
}
|
||
}
|
||
|
||
// QwenVL 调 local-ai 的多模态 VLM（qwen3.6-35b-a3b + mmproj，llama.cpp 后端）做图片推理。
|
||
// 注意：VL 与 z-image 生成不能同时驻留 12G 显存（OOM），调用方须保证阶段互斥。
|
||
func QwenVL(ctx context.Context, imgData []byte, mime, prompt string) (string, error) {
|
||
base := strings.TrimRight(g.Cfg().MustGet(ctx, "localAi.baseUrl").String(), "/")
|
||
if base == "" {
|
||
return "", fmt.Errorf("localAi.baseUrl 未配置")
|
||
}
|
||
model := g.Cfg().MustGet(ctx, "localAi.vlmModel", "qwen3.6-35b-a3b").String()
|
||
b64 := base64.StdEncoding.EncodeToString(imgData)
|
||
body, err := json.Marshal(map[string]any{
|
||
"model": model,
|
||
"messages": []map[string]any{{
|
||
"role": "user",
|
||
"content": []map[string]any{
|
||
{"type": "image_url", "image_url": map[string]string{"url": fmt.Sprintf("data:%s;base64,%s", mime, b64)}},
|
||
{"type": "text", "text": prompt},
|
||
},
|
||
}},
|
||
"max_tokens": 600,
|
||
"temperature": 0.3,
|
||
})
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, base+"/v1/chat/completions", bytes.NewReader(body))
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
req.Header.Set("Content-Type", "application/json")
|
||
resp, err := http.DefaultClient.Do(req)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
defer func() { _ = resp.Body.Close() }()
|
||
raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
if resp.StatusCode != http.StatusOK {
|
||
return "", fmt.Errorf("VLM HTTP %d: %s", resp.StatusCode, truncateStr(string(raw), 200))
|
||
}
|
||
var out struct {
|
||
Choices []struct {
|
||
Message struct {
|
||
Content string `json:"content"`
|
||
} `json:"message"`
|
||
} `json:"choices"`
|
||
}
|
||
if err := json.Unmarshal(raw, &out); err != nil {
|
||
return "", err
|
||
}
|
||
if len(out.Choices) == 0 {
|
||
return "", fmt.Errorf("VLM 无响应内容")
|
||
}
|
||
return out.Choices[0].Message.Content, nil
|
||
}
|
||
|
||
// Detect 对单张图片做全图检测（不做任何裁剪，位置由模型自行推理）：
|
||
// 整图等比缩放至最长边 InputSize 提交，坐标映射回原图像素后返回。
|
||
// imgW/imgH 为原图尺寸；返回坐标均为原图像素尺度。
|
||
func (c *LocalAi) Detect(ctx context.Context, data []byte, mime string, imgW, imgH int) ([]*Detection, error) {
|
||
sub, scale := c.prepare(data, mime, imgW, imgH)
|
||
body, err := json.Marshal(map[string]any{
|
||
"model": c.Model,
|
||
"image": fmt.Sprintf("data:%s;base64,%s", mime, base64.StdEncoding.EncodeToString(sub)),
|
||
"threshold": c.Threshold,
|
||
})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseUrl+"/v1/detection",
|
||
bytes.NewReader(body))
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
req.Header.Set("Content-Type", "application/json")
|
||
resp, err := http.DefaultClient.Do(req)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer func() { _ = resp.Body.Close() }()
|
||
raw, err := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if resp.StatusCode != http.StatusOK {
|
||
return nil, fmt.Errorf("RF-DETR HTTP %d: %s", resp.StatusCode, truncateStr(string(raw), 200))
|
||
}
|
||
var out struct {
|
||
Detections []*Detection `json:"detections"`
|
||
}
|
||
if err := json.Unmarshal(raw, &out); err != nil {
|
||
return nil, err
|
||
}
|
||
// 响应 x/y 是框中心：先转左上角，再从提交图比例尺映射回原图像素。
|
||
// 不做此转换会把中心当左上角，下游再加 w/2 时整框偏移半宽半高
|
||
// （实测 RNPHE_133: 响应(403,664) 即目标中心，旧逻辑落库 cx 偏 +w/2）。
|
||
for _, d := range out.Detections {
|
||
d.X -= d.Width / 2
|
||
d.Y -= d.Height / 2
|
||
d.X /= scale
|
||
d.Y /= scale
|
||
d.Width /= scale
|
||
d.Height /= scale
|
||
}
|
||
return out.Detections, nil
|
||
}
|
||
|
||
// prepare 整图等比缩放至最长边 InputSize（等比，不裁剪），返回提交字节与缩放比（原图/提交图）。
|
||
func (c *LocalAi) prepare(data []byte, mime string, imgW, imgH int) ([]byte, float64) {
|
||
if imgW <= 0 || imgH <= 0 || imgW <= c.InputSize && imgH <= c.InputSize {
|
||
return data, 1
|
||
}
|
||
scale := float64(c.InputSize) / float64(maxInt(imgW, imgH))
|
||
w, h := maxInt(1, int(float64(imgW)*scale)), maxInt(1, int(float64(imgH)*scale))
|
||
src, _, err := image.Decode(bytes.NewReader(data))
|
||
if err != nil {
|
||
return data, 1
|
||
}
|
||
dst := bilinearResize(src, w, h)
|
||
var buf bytes.Buffer
|
||
if strings.Contains(mime, "png") {
|
||
_ = png.Encode(&buf, dst)
|
||
} else {
|
||
_ = jpeg.Encode(&buf, dst, &jpeg.Options{Quality: 92})
|
||
}
|
||
return buf.Bytes(), scale
|
||
}
|
||
|
||
// bilinearResize 双线性缩放（RF-DETR 对小图鲁棒，检测场景无需高质量插值）
|
||
func bilinearResize(src image.Image, w, h int) *image.RGBA {
|
||
b := src.Bounds()
|
||
dst := image.NewRGBA(image.Rect(0, 0, w, h))
|
||
if b.Dx() == 0 || b.Dy() == 0 {
|
||
return dst
|
||
}
|
||
for y := 0; y < h; y++ {
|
||
sy := float64(y) * float64(b.Dy()-1) / float64(maxInt(h-1, 1))
|
||
y0, y1 := int(sy), minInt(int(sy)+1, b.Dy()-1)
|
||
fy := sy - float64(y0)
|
||
for x := 0; x < w; x++ {
|
||
sx := float64(x) * float64(b.Dx()-1) / float64(maxInt(w-1, 1))
|
||
x0, x1 := int(sx), minInt(int(sx)+1, b.Dx()-1)
|
||
fx := sx - float64(x0)
|
||
r00, g00, b00, _ := src.At(b.Min.X+x0, b.Min.Y+y0).RGBA()
|
||
r10, g10, b10, _ := src.At(b.Min.X+x1, b.Min.Y+y0).RGBA()
|
||
r01, g01, b01, _ := src.At(b.Min.X+x0, b.Min.Y+y1).RGBA()
|
||
r11, g11, b11, _ := src.At(b.Min.X+x1, b.Min.Y+y1).RGBA()
|
||
top := func(v00, v10 uint32) uint8 {
|
||
return uint8((float64(v00)*(1-fx) + float64(v10)*fx) / 257)
|
||
}
|
||
bot := func(v01, v11 uint32) uint8 {
|
||
return uint8((float64(v01)*(1-fx) + float64(v11)*fx) / 257)
|
||
}
|
||
r := uint8((float64(top(r00, r10))*(1-fy) + float64(bot(r01, r11))*fy))
|
||
gx := uint8((float64(top(g00, g10))*(1-fy) + float64(bot(g01, g11))*fy))
|
||
bb := uint8((float64(top(b00, b10))*(1-fy) + float64(bot(b01, b11))*fy))
|
||
dst.Set(x, y, color.RGBA{R: r, G: gx, B: bb, A: 255})
|
||
}
|
||
}
|
||
return dst
|
||
}
|
||
|
||
func maxInt(a, b int) int {
|
||
if a > b {
|
||
return a
|
||
}
|
||
return b
|
||
}
|
||
|
||
func minInt(a, b int) int {
|
||
if a < b {
|
||
return a
|
||
}
|
||
return b
|
||
}
|