1

2026-08-26 18:15:54 +08:00
parent 54d343b739
commit a4568d8a55
79 changed files with 11264 additions and 560 deletions
@@ -0,0 +1,176 @@
+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"}]}，坐标单位 = 提交图片像素。
+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 单目标检测结果（提交图比例尺下的像素坐标）
+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.08).Float64(),
+		ConfConfirmed:    g.Cfg().MustGet(ctx, "localAi.confConfirmed", 0.2).Float64(),
+		OverlapThreshold: g.Cfg().MustGet(ctx, "localAi.overlapThreshold", 0.3).Float64(),
+		InputSize:        g.Cfg().MustGet(ctx, "localAi.inputSize", 700).Int(),
+	}
+}
+
+// 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
+	}
+	// 坐标从提交图比例尺映射回原图
+	for _, d := range out.Detections {
+		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
+}