observer/server/common/localai.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"}]}，坐标单位 = 提交图片像素。
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
}