693 lines
23 KiB
Go
693 lines
23 KiB
Go
package service
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"image"
|
||
_ "image/jpeg"
|
||
_ "image/png"
|
||
"math"
|
||
"os"
|
||
"path/filepath"
|
||
"regexp"
|
||
"sort"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/gogf/gf/v2/errors/gerror"
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
"github.com/gogf/gf/v2/os/gtime"
|
||
|
||
"observer-server/biz/consts"
|
||
"observer-server/biz/dao"
|
||
"observer-server/biz/model/dto"
|
||
"observer-server/biz/model/entity"
|
||
"observer-server/common"
|
||
)
|
||
|
||
// labelTaskService 预标注工作台后端：RF-DETR 全图扫描，AI 标注结果直写
|
||
// dataset_image.labels_json（与人工标注同存同编辑），人工可修改/清理全部框；
|
||
// 训练前整理 YOLO 训练集。
|
||
// 全图扫描规则（项目既定）：不套用生成规格的位置裁剪，候选宁多勿漏。
|
||
type labelTaskService struct{}
|
||
|
||
var LabelTask = &labelTaskService{}
|
||
|
||
// recoverLabelTasks 服务重启恢复：孤儿 running 预标注任务置 done + 错误提示
|
||
// （RF-DETR 检测无状态，重新发起即可重新生成标注）。
|
||
func (s *labelTaskService) recoverLabelTasks(ctx context.Context) {
|
||
list, err := dao.LabelTask.ListRunning(ctx)
|
||
if err != nil {
|
||
g.Log().Errorf(ctx, "恢复预标注任务失败: %+v", err)
|
||
return
|
||
}
|
||
for _, t := range list {
|
||
if err := dao.LabelTask.Finish(ctx, t.Id, "服务重启，任务中断，可重新发起"); err != nil {
|
||
g.Log().Errorf(ctx, "恢复预标注任务 %d 失败: %+v", t.Id, err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// AdminListLabelTasks 预标注任务分页（组装数据集名）
|
||
func (s *labelTaskService) AdminListLabelTasks(ctx context.Context, req *dto.AdminLabelTaskListReq) (*dto.AdminLabelTaskListRes, error) {
|
||
page, size := common.NormalizePage(req.Page, req.Size)
|
||
list, total, err := dao.LabelTask.Page(ctx, page, size)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
names := Training.datasetNameMap(ctx)
|
||
items := make([]*dto.AdminLabelTaskItem, 0, len(list))
|
||
for _, v := range list {
|
||
items = append(items, &dto.AdminLabelTaskItem{
|
||
Id: v.Id,
|
||
DatasetId: v.DatasetId,
|
||
DatasetName: names[v.DatasetId],
|
||
Status: v.Status,
|
||
Total: v.Total,
|
||
Done: v.Done,
|
||
Error: v.Error,
|
||
CreatedAt: v.CreatedAt,
|
||
FinishedAt: v.FinishedAt,
|
||
})
|
||
}
|
||
return &dto.AdminLabelTaskListRes{Total: total, List: items}, nil
|
||
}
|
||
|
||
// AdminStartLabelTask 发起预标注：串行检查（数据集存在 + 无 running 任务）→ 插任务 →
|
||
// 池内逐张调 RF-DETR（全图扫描，AI 端点/模型取 config.yml localAi 节点），
|
||
// 完成后标注直写 dataset_image.labels_json 置 done；任一图片失败则任务置 done + error。
|
||
// 多选批量：req.Filenames 非空时只扫选中图。
|
||
func (s *labelTaskService) AdminStartLabelTask(ctx context.Context, req *dto.AdminLabelTaskStartReq) (*dto.AdminLabelTaskStartRes, error) {
|
||
dataset, err := dao.Dataset.GetById(ctx, req.DatasetId)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if dataset == nil {
|
||
return nil, gerror.NewCode(common.CodeDatasetNotFound)
|
||
}
|
||
if common.LocalAiClient(ctx) == nil {
|
||
return nil, gerror.NewCode(common.CodeLocalAiNotConfigured)
|
||
}
|
||
images, err := dao.DatasetImage.ListByDataset(ctx, dataset.Id)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if len(images) == 0 {
|
||
return nil, gerror.New("数据集暂无图片")
|
||
}
|
||
// 多选批量：只处理选中图片（缺省全量）
|
||
subset := len(req.Filenames) > 0
|
||
filenamesJSON := ""
|
||
if subset {
|
||
sel := make(map[string]bool, len(req.Filenames))
|
||
for _, f := range req.Filenames {
|
||
sel[f] = true
|
||
}
|
||
filtered := make([]*entity.DatasetImage, 0, len(sel))
|
||
for _, img := range images {
|
||
if sel[img.Filename] {
|
||
filtered = append(filtered, img)
|
||
}
|
||
}
|
||
if len(filtered) == 0 {
|
||
return nil, gerror.New("选中的图片不在该数据集内")
|
||
}
|
||
images = filtered
|
||
names := make([]string, 0, len(images))
|
||
for _, img := range images {
|
||
names = append(names, img.Filename)
|
||
}
|
||
if raw, jErr := json.Marshal(names); jErr == nil {
|
||
filenamesJSON = string(raw)
|
||
}
|
||
}
|
||
taskId, err := s.startDetection(ctx, dataset, images, filenamesJSON)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return &dto.AdminLabelTaskStartRes{Id: taskId}, nil
|
||
}
|
||
|
||
// startDetection 发起预标注任务：Serial 内检查并发（已有 running 任务报错）→ 插任务 → 启动检测协程
|
||
func (s *labelTaskService) startDetection(ctx context.Context, dataset *entity.Dataset, images []*entity.DatasetImage, filenamesJSON string) (int64, error) {
|
||
now := gtime.Now()
|
||
var taskId int64
|
||
err := common.Serial().Submit(ctx, func() error {
|
||
running, err := dao.LabelTask.GetRunningByDataset(ctx, dataset.Id)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if running != nil {
|
||
return gerror.NewCode(common.CodeLabelTaskRunning)
|
||
}
|
||
taskId, err = dao.LabelTask.Insert(ctx, &entity.LabelTask{
|
||
DatasetId: dataset.Id,
|
||
Status: consts.LabelTaskRunning,
|
||
Total: len(images),
|
||
Filenames: filenamesJSON,
|
||
CreatedAt: now,
|
||
})
|
||
return err
|
||
})
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
s.runDetection(ctx, taskId, dataset, images, common.LocalAiClient(ctx))
|
||
return taskId, nil
|
||
}
|
||
|
||
// AutoLabel 图片入库自动触发标注（上传/AI 生成成功后调用）：只标本次新增图。
|
||
// localAi 未配置 → 报错（调用方拒绝/回滚，不允许产生无标注图）；
|
||
// 已有 running 任务（忙）→ 不报错：本次跳过，由任务成功完成后的自动补标轮兜底。
|
||
func (s *labelTaskService) AutoLabel(ctx context.Context, dataset *entity.Dataset, images []*entity.DatasetImage) error {
|
||
if len(images) == 0 {
|
||
return nil
|
||
}
|
||
if common.LocalAiClient(ctx) == nil {
|
||
return gerror.NewCode(common.CodeLocalAiNotConfigured)
|
||
}
|
||
names := make([]string, 0, len(images))
|
||
for _, img := range images {
|
||
names = append(names, img.Filename)
|
||
}
|
||
filenamesJSON := ""
|
||
if raw, jErr := json.Marshal(names); jErr == nil {
|
||
filenamesJSON = string(raw)
|
||
}
|
||
_, err := s.startDetection(ctx, dataset, images, filenamesJSON)
|
||
if err != nil && gerror.HasCode(err, common.CodeLabelTaskRunning) {
|
||
return nil // 忙：不报错
|
||
}
|
||
return err
|
||
}
|
||
|
||
// autoSupplement 自动补标：任务成功完成后，该数据集仍有未标注图（忙时入库的图/人工清空的图）
|
||
// 则续一轮只标未标注图；失败任务不续（防配置故障时无限重试）。
|
||
func (s *labelTaskService) autoSupplement(ctx context.Context, dataset *entity.Dataset) {
|
||
unlabeled, err := dao.DatasetImage.ListUnlabeledByDataset(ctx, dataset.Id)
|
||
if err != nil || len(unlabeled) == 0 {
|
||
return
|
||
}
|
||
if err := s.AutoLabel(ctx, dataset, unlabeled); err != nil {
|
||
g.Log().Infof(ctx, "自动补标未触发（数据集 %d）: %v", dataset.Id, err)
|
||
}
|
||
}
|
||
|
||
// runDetection 预标注执行协程（生命周期任务）：池内逐张检测，进度经 Serial 更新。
|
||
// 每张：读图 → 全图检测（等比缩放提交，坐标映射回原图归一化）→ 标注（conf≥confConfirmed 为 class 0）。
|
||
// 全部成功后 Serial 内逐张覆写 labels_json（重跑覆盖该图标注）。
|
||
func (s *labelTaskService) runDetection(ctx context.Context, taskId int64, dataset *entity.Dataset, images []*entity.DatasetImage, client *common.LocalAi) {
|
||
// 生命周期任务必须脱离请求 ctx：请求结束即取消，会让 Submit 秒退 + Finish 静默失败 → 任务悬挂
|
||
bgCtx := context.Background()
|
||
go func() {
|
||
if client == nil {
|
||
_ = dao.LabelTask.Finish(bgCtx, taskId, "标注服务未配置")
|
||
return
|
||
}
|
||
dir := common.DatasetImagesDir(bgCtx, dataset.Name)
|
||
results := make([][]*dto.AdminLabelBox, len(images))
|
||
failed := ""
|
||
doneCount := 0
|
||
for i, img := range images {
|
||
// 单张 120s 超时兜底（无超时 + 无取消的独立 ctx 下防检测服务悬挂拖死任务）
|
||
detCtx, cancel := context.WithTimeout(bgCtx, 120*time.Second)
|
||
err := common.LabelTaskPoolInstance().Submit(detCtx, func(ctx context.Context) error {
|
||
data, err := os.ReadFile(filepath.Join(dir, img.Filename))
|
||
if err != nil {
|
||
return gerror.Wrapf(err, "读取图片失败")
|
||
}
|
||
w, h := imageSize(data)
|
||
if w <= 0 || h <= 0 {
|
||
return gerror.New("无法识别图片尺寸")
|
||
}
|
||
mime := imageMime(img.Filename)
|
||
detections, err := client.Detect(ctx, data, mime, w, h)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
cands := make([]*dto.AdminLabelBox, 0, len(detections))
|
||
for _, d := range detections {
|
||
// 坐标映射回原图像素后归一化，越界轻微裁剪
|
||
box := &dto.AdminLabelBox{
|
||
Cx: clamp01((d.X + d.Width/2) / float64(w)),
|
||
Cy: clamp01((d.Y + d.Height/2) / float64(h)),
|
||
W: clamp01(d.Width / float64(w)),
|
||
H: clamp01(d.Height / float64(h)),
|
||
Confidence: d.Confidence,
|
||
Class: 1,
|
||
}
|
||
if d.Confidence >= client.ConfConfirmed {
|
||
box.Class = 0
|
||
}
|
||
cands = append(cands, box)
|
||
}
|
||
boxes := suppressOverlap(cands, client.OverlapThreshold)
|
||
// 生成时已指定目标数量上限：自动标注按置信度裁剪到该数
|
||
// （生成校验已保证图内目标 ≤N，宁多勿漏阈值带出的低置信噪声框在此去除）
|
||
if img.AnimalCount > 0 && len(boxes) > img.AnimalCount {
|
||
sort.Slice(boxes, func(a, b int) bool { return boxes[a].Confidence > boxes[b].Confidence })
|
||
boxes = boxes[:img.AnimalCount]
|
||
}
|
||
results[i] = boxes
|
||
return nil
|
||
})
|
||
cancel()
|
||
if err != nil {
|
||
failed = fmt.Sprintf("第 %d 张(%s)检测失败: %v", doneCount+1, img.Filename, err)
|
||
break
|
||
}
|
||
doneCount++
|
||
_ = common.Serial().Submit(bgCtx, func() error {
|
||
return dao.LabelTask.UpdateProgress(bgCtx, taskId, doneCount)
|
||
})
|
||
}
|
||
if failed != "" {
|
||
_ = dao.LabelTask.Finish(bgCtx, taskId, failed)
|
||
return
|
||
}
|
||
// 全部成功：Serial 内逐张覆写标注（重跑覆盖该图标注，人工保存同走 UpdateLabels）
|
||
err := common.Serial().Submit(bgCtx, func() error {
|
||
for i := range results {
|
||
raw, jErr := json.Marshal(results[i])
|
||
if jErr != nil {
|
||
return gerror.Wrap(jErr, "标注序列化失败")
|
||
}
|
||
if uErr := dao.DatasetImage.UpdateLabels(bgCtx, images[i].Id, string(raw)); uErr != nil {
|
||
return uErr
|
||
}
|
||
}
|
||
return nil
|
||
})
|
||
if err != nil {
|
||
_ = dao.LabelTask.Finish(bgCtx, taskId, err.Error())
|
||
return
|
||
}
|
||
_ = dao.LabelTask.Finish(bgCtx, taskId, "")
|
||
// 成功完成：自动补标该数据集未标注图（忙时入库的图等），失败任务不续
|
||
s.autoSupplement(bgCtx, dataset)
|
||
}()
|
||
}
|
||
|
||
// AdminImageVlmReview 单图 VLM 藏匿位补检（两阶段标注第二阶段，须与 z-image 生成显存互斥）：
|
||
// 以当前 labels_json（RF-DETR/人工框）为排除集，调 qwen3.8-9b(+mmproj) 按环境/光线/习性
|
||
// 推理可能藏匿目标的位置，取前 3 个追加为疑似框(class 1)。
|
||
func (s *labelTaskService) AdminImageVlmReview(ctx context.Context, req *dto.AdminImageVlmReviewReq) (*dto.AdminImageVlmReviewRes, error) {
|
||
dataset, err := dao.Dataset.GetById(ctx, req.DatasetId)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if dataset == nil {
|
||
return nil, gerror.NewCode(common.CodeDatasetNotFound)
|
||
}
|
||
img, err := dao.DatasetImage.GetById(ctx, req.ImageId)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if img == nil || img.DatasetId != dataset.Id {
|
||
return nil, gerror.NewCode(common.CodeImageNotFound)
|
||
}
|
||
data, err := os.ReadFile(filepath.Join(common.DatasetImagesDir(ctx, dataset.Name), img.Filename))
|
||
if err != nil {
|
||
return nil, gerror.Wrap(err, "读取图片失败")
|
||
}
|
||
var existing []*dto.AdminLabelBox
|
||
if strings.TrimSpace(img.LabelsJson) != "" {
|
||
if jErr := json.Unmarshal([]byte(img.LabelsJson), &existing); jErr != nil {
|
||
return nil, gerror.Wrap(jErr, "解析已有标注失败")
|
||
}
|
||
}
|
||
// 排除集：已确认框坐标（VLM 不得重复输出这些位置）
|
||
exclude := "无"
|
||
if len(existing) > 0 {
|
||
parts := make([]string, 0, len(existing))
|
||
for _, b := range existing {
|
||
parts = append(parts, fmt.Sprintf("(%.3f,%.3f,%.3f,%.3f)", b.Cx, b.Cy, b.W, b.H))
|
||
}
|
||
exclude = strings.Join(parts, " ")
|
||
}
|
||
// 物种：从生成提示词与数据集物种池匹配（匹配不到交给 VLM 从已确认框自行判断）
|
||
species := ""
|
||
if pool := g.Cfg().MustGet(ctx, "imageGen.speciesByDataset."+dataset.Name).Strings(); len(pool) > 0 {
|
||
for _, name := range pool {
|
||
if strings.Contains(img.Prompt, name) {
|
||
species = name
|
||
break
|
||
}
|
||
}
|
||
}
|
||
speciesLine := "从已确认目标框自行判断物种"
|
||
if species != "" {
|
||
speciesLine = species
|
||
}
|
||
prompt := fmt.Sprintf(
|
||
"你是野生动物监测照片分析助手。\n物种：%s\n拍摄场景与时间光线线索：%s\n已确认目标位置（归一化 cx,cy,w,h，已排除，不要再输出这些位置）：%s\n"+
|
||
"请基于环境、天气、光线与该物种习性（觅食路线、隐蔽处、阴影边缘、植被深处），分析除已确认位置外还可能藏匿个体的位置。\n"+
|
||
"输出归一化 bbox 的 JSON 数组：[{\"cx\":0.5,\"cy\":0.5,\"w\":0.2,\"h\":0.15}]，最多 3 个，按可能性从高到低；确无可能则输出 []。只输出 JSON 数组本身，不要其他文字。",
|
||
speciesLine, img.Prompt, exclude)
|
||
content, err := common.QwenVL(ctx, data, imageMime(img.Filename), prompt)
|
||
if err != nil {
|
||
return nil, gerror.Wrap(err, "VLM 推理失败")
|
||
}
|
||
// 提取首个 JSON 数组并校验
|
||
m := regexp.MustCompile(`\[[\s\S]*?\]`).FindString(content)
|
||
res := &dto.AdminImageVlmReviewRes{}
|
||
if m == "" {
|
||
res.Note = "VLM 未输出坐标数组"
|
||
return res, nil
|
||
}
|
||
var raw []struct {
|
||
Cx, Cy, W, H float64
|
||
}
|
||
if err := json.Unmarshal([]byte(m), &raw); err != nil {
|
||
res.Note = "VLM 坐标解析失败"
|
||
return res, nil
|
||
}
|
||
client := common.LocalAiClient(ctx)
|
||
overlap := 0.3
|
||
if client != nil {
|
||
overlap = client.OverlapThreshold
|
||
}
|
||
for _, b := range raw {
|
||
if len(res.Boxes) >= 3 {
|
||
break
|
||
}
|
||
if b.Cx <= 0 || b.Cx >= 1 || b.Cy <= 0 || b.Cy >= 1 || b.W <= 0 || b.W >= 1 || b.H <= 0 || b.H >= 1 {
|
||
continue
|
||
}
|
||
nb := &dto.AdminLabelBox{Cx: b.Cx, Cy: b.Cy, W: b.W, H: b.H, Confidence: 0.1, Class: 1}
|
||
dup := false
|
||
for _, e := range existing {
|
||
if boxOverlap(nb, e) > overlap {
|
||
dup = true
|
||
break
|
||
}
|
||
}
|
||
if dup {
|
||
continue
|
||
}
|
||
res.Boxes = append(res.Boxes, nb)
|
||
}
|
||
if len(res.Boxes) == 0 {
|
||
res.Note = "VLM 未给出有效的藏匿位坐标"
|
||
return res, nil
|
||
}
|
||
// 追加为疑似框并落库（与 RF-DETR/人工框同层，工作台可编辑）
|
||
appended := append(existing, res.Boxes...)
|
||
if err := common.Serial().Submit(ctx, func() error {
|
||
rawJson, jErr := json.Marshal(appended)
|
||
if jErr != nil {
|
||
return jErr
|
||
}
|
||
return dao.DatasetImage.UpdateLabels(ctx, img.Id, string(rawJson))
|
||
}); err != nil {
|
||
return nil, err
|
||
}
|
||
res.Added = len(res.Boxes)
|
||
return res, nil
|
||
}
|
||
|
||
// AdminLabelTaskDetail 预标注任务详情：标注（labels_json）输出，供工作台 canvas 叠框；
|
||
// 图片尺寸读取文件头。
|
||
func (s *labelTaskService) AdminLabelTaskDetail(ctx context.Context, req *dto.AdminLabelTaskDetailReq) (*dto.AdminLabelTaskDetailRes, error) {
|
||
t, err := dao.LabelTask.GetById(ctx, req.Id)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if t == nil {
|
||
return nil, gerror.New("标注任务不存在")
|
||
}
|
||
dataset, err := dao.Dataset.GetById(ctx, t.DatasetId)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if dataset == nil {
|
||
return nil, gerror.NewCode(common.CodeDatasetNotFound)
|
||
}
|
||
items, err := s.buildWorkbenchItems(ctx, dataset)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return &dto.AdminLabelTaskDetailRes{
|
||
TaskId: t.Id,
|
||
DatasetId: t.DatasetId,
|
||
DatasetName: dataset.Name,
|
||
Status: t.Status,
|
||
Total: t.Total,
|
||
Done: t.Done,
|
||
Error: t.Error,
|
||
Images: items,
|
||
}, nil
|
||
}
|
||
|
||
// buildWorkbenchItems 工作台单张图数据：全部标注（labels_json，AI 自动标注与人工框同层）；
|
||
// 尺寸统一读文件头。
|
||
func (s *labelTaskService) buildWorkbenchItems(ctx context.Context, dataset *entity.Dataset) ([]*dto.AdminLabelImageItem, error) {
|
||
images, err := dao.DatasetImage.ListByDataset(ctx, dataset.Id)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
imgDir := common.DatasetImagesDir(ctx, dataset.Name)
|
||
items := make([]*dto.AdminLabelImageItem, 0, len(images))
|
||
for _, img := range images {
|
||
item := &dto.AdminLabelImageItem{
|
||
Filename: img.Filename,
|
||
Url: datasetImageUrl(ctx, dataset.Id, img.Filename),
|
||
}
|
||
if img.LabelsJson != "" && img.LabelsJson != "[]" {
|
||
var boxes []*dto.AdminLabelBox
|
||
if json.Unmarshal([]byte(img.LabelsJson), &boxes) == nil && len(boxes) > 0 {
|
||
item.Boxes = boxes
|
||
item.Labeled = true
|
||
}
|
||
}
|
||
if data, rErr := os.ReadFile(filepath.Join(imgDir, img.Filename)); rErr == nil {
|
||
item.Width, item.Height = imageSize(data)
|
||
}
|
||
items = append(items, item)
|
||
}
|
||
return items, nil
|
||
}
|
||
|
||
// AdminLabelWorkbench 标注工作台数据：数据集无历史任务时直接输出图片 + 全部标注，
|
||
// 标注读 dataset_image.labels_json（与任务详情同一组装）。
|
||
func (s *labelTaskService) AdminLabelWorkbench(ctx context.Context, req *dto.AdminLabelWorkbenchReq) (*dto.AdminLabelWorkbenchRes, error) {
|
||
dataset, err := dao.Dataset.GetById(ctx, req.DatasetId)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if dataset == nil {
|
||
return nil, gerror.NewCode(common.CodeDatasetNotFound)
|
||
}
|
||
items, err := s.buildWorkbenchItems(ctx, dataset)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return &dto.AdminLabelWorkbenchRes{
|
||
DatasetId: dataset.Id,
|
||
DatasetName: dataset.Name,
|
||
Images: items,
|
||
}, nil
|
||
}
|
||
|
||
// AdminLabelSave 保存单张图标注：校验坐标 → json.Marshal 覆写 labels_json（空框=清空），
|
||
// 并刷新数据集 labeled_count（labels_json 非空数组的图片数）。
|
||
func (s *labelTaskService) AdminLabelSave(ctx context.Context, req *dto.AdminLabelSaveReq) (*dto.AdminLabelSaveRes, error) {
|
||
dataset, err := dao.Dataset.GetById(ctx, req.DatasetId)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if dataset == nil {
|
||
return nil, gerror.NewCode(common.CodeDatasetNotFound)
|
||
}
|
||
img, err := dao.DatasetImage.GetByFilename(ctx, dataset.Id, req.Filename)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if img == nil {
|
||
return nil, gerror.NewCode(common.CodeImageNotFound)
|
||
}
|
||
for _, box := range req.Boxes {
|
||
if box.Cx < 0 || box.Cy < 0 || box.W <= 0 || box.H <= 0 || box.Cx > 1 || box.Cy > 1 {
|
||
return nil, gerror.New("标注框坐标非法（需 0~1 归一化）")
|
||
}
|
||
}
|
||
// 空框 = 清空标注
|
||
raw := ""
|
||
if len(req.Boxes) > 0 {
|
||
b, err := json.Marshal(req.Boxes)
|
||
if err != nil {
|
||
return nil, gerror.Wrap(err, "标注序列化失败")
|
||
}
|
||
raw = string(b)
|
||
}
|
||
var labeled int64
|
||
if err := common.Serial().Submit(ctx, func() error {
|
||
if err := dao.DatasetImage.UpdateLabels(ctx, img.Id, raw); err != nil {
|
||
return err
|
||
}
|
||
labeled, err = dao.DatasetImage.CountLabeledByDataset(ctx, dataset.Id)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
status := ""
|
||
if labeled > 0 {
|
||
status = consts.DatasetStatusLabeled
|
||
}
|
||
return dao.Dataset.UpdateCounters(ctx, dataset.Id, 0, labeled, status)
|
||
}); err != nil {
|
||
return nil, err
|
||
}
|
||
return &dto.AdminLabelSaveRes{LabeledCount: labeled}, nil
|
||
}
|
||
|
||
// prepareYoloSet 训练前组装内存 YOLO 训练集包：已标注图（labels_json 非空）按 80/20 拆 train/val，
|
||
// 标注 txt 内存生成、原图仅记源路径（由训练通道读取，不落本地暂存盘）；无标注报错。
|
||
// data.yaml 由训练发起方追加进包（path 需指向训练机）。
|
||
func (s *labelTaskService) prepareYoloSet(ctx context.Context, dataset *entity.Dataset) (*common.YoloPackage, error) {
|
||
images, err := dao.DatasetImage.ListByDataset(ctx, dataset.Id)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
type item struct {
|
||
filename string
|
||
lines string
|
||
}
|
||
var kept []item
|
||
for _, img := range images {
|
||
if img.LabelsJson == "" || img.LabelsJson == "[]" {
|
||
continue // 空标注图剔除
|
||
}
|
||
var boxes []*dto.AdminLabelBox
|
||
if json.Unmarshal([]byte(img.LabelsJson), &boxes) != nil || len(boxes) == 0 {
|
||
continue
|
||
}
|
||
var b strings.Builder
|
||
for _, box := range boxes {
|
||
fmt.Fprintf(&b, "%d %.6f %.6f %.6f %.6f\n", box.Class, box.Cx, box.Cy, box.W, box.H)
|
||
}
|
||
kept = append(kept, item{filename: img.Filename, lines: strings.TrimSpace(b.String())})
|
||
}
|
||
if len(kept) == 0 {
|
||
return nil, gerror.New("数据集无有效标注，请先在标注工作台完成标注")
|
||
}
|
||
// 固定随机种子 + 20% val（至少 1 张，语义沿用原 prepare_yolo.py）
|
||
idx := make([]int, len(kept))
|
||
for i := range idx {
|
||
idx[i] = i
|
||
}
|
||
randShuffle(idx)
|
||
nVal := len(kept) / 5
|
||
if nVal < 1 {
|
||
nVal = 1
|
||
}
|
||
imgSrc := common.DatasetImagesDir(ctx, dataset.Name)
|
||
pkg := &common.YoloPackage{}
|
||
addSplit := func(split string, items []item) {
|
||
for _, it := range items {
|
||
pkg.Files = append(pkg.Files,
|
||
common.YoloFile{Name: filepath.Join("images", split, it.filename), ImagePath: filepath.Join(imgSrc, it.filename)},
|
||
common.YoloFile{Name: filepath.Join("labels", split, strings.TrimSuffix(it.filename, filepath.Ext(it.filename))+".txt"), Content: []byte(it.lines + "\n")},
|
||
)
|
||
}
|
||
}
|
||
var trainItems, valItems []item
|
||
for i, it := range kept {
|
||
if i < nVal {
|
||
valItems = append(valItems, it)
|
||
} else {
|
||
trainItems = append(trainItems, it)
|
||
}
|
||
}
|
||
addSplit("train", trainItems)
|
||
addSplit("val", valItems)
|
||
return pkg, nil
|
||
}
|
||
|
||
// randShuffle Fisher-Yates 伪随机（固定种子，沿用原 prepare_yolo.py random.seed(42) 语义）
|
||
func randShuffle(n []int) {
|
||
state := uint32(42)
|
||
seed := func() uint32 {
|
||
state = state*1664525 + 1013904223
|
||
return state
|
||
}
|
||
for i := len(n) - 1; i > 0; i-- {
|
||
j := int(seed() % uint32(i+1))
|
||
n[i], n[j] = n[j], n[i]
|
||
}
|
||
}
|
||
|
||
// suppressOverlap 重叠去重（NMS 风格）：按置信度降序依次保留，与已保留框重叠比 > overlapThreshold
|
||
// 的框剔除（RF-DETR 同一目标重复检出时多个高度重叠框，只留置信度最高者；跨 class 去重）。
|
||
func suppressOverlap(boxes []*dto.AdminLabelBox, overlapThreshold float64) []*dto.AdminLabelBox {
|
||
if len(boxes) <= 1 {
|
||
return boxes
|
||
}
|
||
sorted := append([]*dto.AdminLabelBox(nil), boxes...)
|
||
sort.Slice(sorted, func(i, j int) bool {
|
||
return sorted[i].Confidence > sorted[j].Confidence
|
||
})
|
||
kept := make([]*dto.AdminLabelBox, 0, len(sorted))
|
||
for i := range sorted {
|
||
dup := false
|
||
for _, k := range kept {
|
||
if boxOverlap(sorted[i], k) > overlapThreshold {
|
||
dup = true
|
||
break
|
||
}
|
||
}
|
||
if !dup {
|
||
kept = append(kept, sorted[i])
|
||
}
|
||
}
|
||
return kept
|
||
}
|
||
|
||
// boxOverlap 两个归一化框（cx,cy,w,h）的重叠比（minIoU）：交叠面积 / 两框较小面积。
|
||
// 用 minIoU 而非 IoU：RF-DETR 对同一目标常输出一大一小两个框（IoU 仅 0.3~0.5），
|
||
// 大框套小框时小框被覆盖比例高（0.3~0.9）能命中；相邻目标两框互有外露，minIoU 通常 < 0.3。
|
||
func boxOverlap(a, b *dto.AdminLabelBox) float64 {
|
||
ax1, ay1, ax2, ay2 := a.Cx-a.W/2, a.Cy-a.H/2, a.Cx+a.W/2, a.Cy+a.H/2
|
||
bx1, by1, bx2, by2 := b.Cx-b.W/2, b.Cy-b.H/2, b.Cx+b.W/2, b.Cy+b.H/2
|
||
ix1, iy1 := math.Max(ax1, bx1), math.Max(ay1, by1)
|
||
ix2, iy2 := math.Min(ax2, bx2), math.Min(ay2, by2)
|
||
if ix2 <= ix1 || iy2 <= iy1 {
|
||
return 0
|
||
}
|
||
inter := (ix2 - ix1) * (iy2 - iy1)
|
||
minArea := math.Min(a.W*a.H, b.W*b.H)
|
||
if minArea <= 0 {
|
||
return 0
|
||
}
|
||
return inter / minArea
|
||
}
|
||
|
||
// imageSize 读取图片尺寸（文件头，不解码全图）
|
||
func imageSize(data []byte) (int, int) {
|
||
cfg, _, err := image.DecodeConfig(bytes.NewReader(data))
|
||
if err != nil {
|
||
return 0, 0
|
||
}
|
||
return cfg.Width, cfg.Height
|
||
}
|
||
|
||
func imageMime(filename string) string {
|
||
switch strings.ToLower(filepath.Ext(filename)) {
|
||
case ".png":
|
||
return "image/png"
|
||
default:
|
||
return "image/jpeg"
|
||
}
|
||
}
|
||
|
||
func clamp01(v float64) float64 {
|
||
if v < 0 {
|
||
return 0
|
||
}
|
||
if v > 1 {
|
||
return 1
|
||
}
|
||
return v
|
||
}
|