336 lines
12 KiB
Go
336 lines
12 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
|
||
// 四级漏斗切片级（技术设计.md「预标注四级漏斗」）：切片检测置信度下限 / 块长边（原图像素）/ 相邻块重叠比
|
||
TileThreshold float64
|
||
TileSize int
|
||
TileOverlap float64
|
||
}
|
||
|
||
// 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(),
|
||
TileThreshold: g.Cfg().MustGet(ctx, "localAi.tileThreshold", 0.12).Float64(),
|
||
TileSize: g.Cfg().MustGet(ctx, "localAi.tileSize", 1024).Int(),
|
||
TileOverlap: g.Cfg().MustGet(ctx, "localAi.tileOverlap", 0.25).Float64(),
|
||
}
|
||
}
|
||
|
||
// QwenVL 调多模态 VLM 做图片推理（OpenAI 兼容 /v1/chat/completions）。
|
||
// 2026-09-10 起默认走 omlx（localAi.vlmBaseUrl，MLX 多模态模型）；未配置 vlmBaseUrl
|
||
// 时回落 localAi.baseUrl（旧 local-ai 部署同机混部形态）。
|
||
func QwenVL(ctx context.Context, imgData []byte, mime, prompt string) (string, error) {
|
||
base := g.Cfg().MustGet(ctx, "localAi.vlmBaseUrl").String()
|
||
if base == "" {
|
||
base = g.Cfg().MustGet(ctx, "localAi.baseUrl").String()
|
||
}
|
||
base = strings.TrimRight(base, "/")
|
||
if base == "" {
|
||
return "", fmt.Errorf("localAi.vlmBaseUrl/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")
|
||
if key := g.Cfg().MustGet(ctx, "localAi.vlmApiKey").String(); key != "" {
|
||
req.Header.Set("Authorization", "Bearer "+key)
|
||
}
|
||
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)
|
||
return c.submitDetect(ctx, sub, mime, scale, 0, 0, c.Threshold)
|
||
}
|
||
|
||
// DetectRegion 对原图中指定像素矩形区域做检测（四级漏斗切片级/VLM 候选区精修）：
|
||
// 裁剪 → 等比缩放至最长边 InputSize 提交 → 坐标映射回原图像素。
|
||
// threshold 由调用方给定（切片级用低于全图的阈值，小目标置信度天然偏低）；
|
||
// region 越界内部钳制到图片边界。返回坐标均为原图像素尺度。
|
||
func (c *LocalAi) DetectRegion(ctx context.Context, data []byte, mime string, imgW, imgH int, rx, ry, rw, rh int, threshold float64) ([]*Detection, error) {
|
||
rx = maxInt(0, minInt(rx, imgW-1))
|
||
ry = maxInt(0, minInt(ry, imgH-1))
|
||
rw = minInt(rw, imgW-rx)
|
||
rh = minInt(rh, imgH-ry)
|
||
if rw <= 0 || rh <= 0 {
|
||
return nil, nil
|
||
}
|
||
src, _, err := image.Decode(bytes.NewReader(data))
|
||
if err != nil {
|
||
return nil, fmt.Errorf("解码图片失败: %w", err)
|
||
}
|
||
cropper, ok := src.(interface {
|
||
SubImage(image.Rectangle) image.Image
|
||
})
|
||
if !ok {
|
||
return nil, fmt.Errorf("图片格式不支持区域裁剪")
|
||
}
|
||
sub := cropper.SubImage(image.Rect(rx, ry, rx+rw, ry+rh))
|
||
var buf bytes.Buffer
|
||
if strings.Contains(mime, "png") {
|
||
_ = png.Encode(&buf, sub)
|
||
} else {
|
||
_ = jpeg.Encode(&buf, sub, &jpeg.Options{Quality: 92})
|
||
}
|
||
region, scale := c.prepare(buf.Bytes(), mime, rw, rh)
|
||
return c.submitDetect(ctx, region, mime, scale, float64(rx), float64(ry), threshold)
|
||
}
|
||
|
||
// submitDetect 提交检测请求并解析：sub 为已缩放的提交图字节，scale 为提交图→原图（或区域）
|
||
// 的缩放比，offX/offY 为提交图坐标系到原图坐标系的偏移（全图 0,0，区域检测为区域左上角），
|
||
// threshold 为本次候选置信度下限。响应 x/y 是框中心：先转左上角，再按比例尺映射回原图像素。
|
||
// 不做中心转左上会把中心当角点，下游再加 w/2 时整框偏移半宽半高
|
||
// （实测 RNPHE_133: 响应(403,664) 即目标中心，旧逻辑落库 cx 偏 +w/2）。
|
||
func (c *LocalAi) submitDetect(ctx context.Context, sub []byte, mime string, scale, offX, offY, threshold float64) ([]*Detection, error) {
|
||
body, err := json.Marshal(map[string]any{
|
||
"model": c.Model,
|
||
"image": fmt.Sprintf("data:%s;base64,%s", mime, base64.StdEncoding.EncodeToString(sub)),
|
||
"threshold": 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
|
||
}
|
||
for _, d := range out.Detections {
|
||
d.X -= d.Width / 2
|
||
d.Y -= d.Height / 2
|
||
d.X = d.X/scale + offX
|
||
d.Y = d.Y/scale + offY
|
||
d.Width /= scale
|
||
d.Height /= scale
|
||
}
|
||
return out.Detections, nil
|
||
}
|
||
|
||
// TileRegion 滑窗切片块（原图像素）
|
||
type TileRegion struct {
|
||
X, Y, W, H int
|
||
}
|
||
|
||
// TileRegions 滑窗切片网格（四级漏斗 L2）：tileSize 为块长边（原图像素），tileOverlap 为
|
||
// 相邻块重叠比（步长 = tileSize*(1-tileOverlap)），从左上到右下枚举，边缘块贴边收口保证全覆盖。
|
||
func TileRegions(imgW, imgH, tileSize int, tileOverlap float64) []TileRegion {
|
||
if tileSize <= 0 {
|
||
tileSize = imgW
|
||
}
|
||
if tileSize > imgW {
|
||
tileSize = imgW
|
||
}
|
||
if tileSize > imgH {
|
||
tileSize = imgH
|
||
}
|
||
step := maxInt(1, int(float64(tileSize)*(1-tileOverlap)))
|
||
tiles := make([]TileRegion, 0, (imgW/step+1)*(imgH/step+1))
|
||
for y := 0; y < imgH; y += step {
|
||
ty := y
|
||
if ty+tileSize > imgH {
|
||
ty = imgH - tileSize
|
||
}
|
||
th := minInt(tileSize, imgH-ty)
|
||
for x := 0; x < imgW; x += step {
|
||
tx := x
|
||
if tx+tileSize > imgW {
|
||
tx = imgW - tileSize
|
||
}
|
||
tw := minInt(tileSize, imgW-tx)
|
||
tiles = append(tiles, TileRegion{X: tx, Y: ty, W: tw, H: th})
|
||
if tx+tw >= imgW {
|
||
break
|
||
}
|
||
}
|
||
if ty+th >= imgH {
|
||
break
|
||
}
|
||
}
|
||
return tiles
|
||
}
|
||
|
||
// 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
|
||
}
|