Files
observer/server/biz/service/dataset.go
T
2026-08-28 14:31:13 +08:00

967 lines
36 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package service
import (
"bytes"
"context"
"fmt"
"image"
"image/jpeg"
"math"
"math/rand"
"net/url"
"os"
"path/filepath"
"regexp"
"strconv"
"strings"
"time"
"github.com/gogf/gf/v2/errors/gerror"
"github.com/gogf/gf/v2/frame/g"
"github.com/gogf/gf/v2/net/ghttp"
"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"
)
// datasetService 数据集业务:管理端「数据集管理」模块。
// 图片文件在 app.datasetDir/datasets/<name>/(平铺、文件名唯一),DB 只存元数据;
// AI 生成为付费资产,删除类接口由前端带确认文案(后端不额外拦截)。
type datasetService struct{}
var Dataset = &datasetService{}
// imgExts 允许上传/保存的图片扩展名
var imgExts = map[string]bool{".jpg": true, ".jpeg": true, ".png": true}
// datasetNameRe 数据集名称规则(与新建接口 DTO 校验一致:中文/字母/数字/下划线/短横线,1-50)
// 更新接口的名称为可选字段,空值合法,正则不能进 DTO 校验链,放 service 侧判定。
var datasetNameRe = regexp.MustCompile(`^[a-zA-Z0-9_一-龥-]{1,50}$`)
// AdminListDatasets 数据集分页列表(卡片展示:封面/描述/训练配置标记 + 最新训练状态聚合)
func (s *datasetService) AdminListDatasets(ctx context.Context, req *dto.AdminDatasetListReq) (*dto.AdminDatasetListRes, error) {
page, size := common.NormalizePage(req.Page, req.Size)
var list []*entity.Dataset
var total int64
var err error
if req.Keyword != "" {
list, total, err = dao.Dataset.PageByKeyword(ctx, req.Keyword, page, size)
} else {
list, total, err = dao.Dataset.Page(ctx, page, size)
}
if err != nil {
return nil, err
}
ids := make([]int64, 0, len(list))
for _, v := range list {
ids = append(ids, v.Id)
}
latest, err := dao.Training.LatestByDatasets(ctx, ids)
if err != nil {
return nil, err
}
// 已发布过的训练不再返回发布按钮(model_version.training_id 反查)
trainingIds := make([]int64, 0, len(latest))
for _, t := range latest {
trainingIds = append(trainingIds, t.Id)
}
published, err := dao.ModelVersion.PublishedByTrainingIds(ctx, trainingIds)
if err != nil {
return nil, err
}
items := make([]*dto.AdminDatasetItem, 0, len(list))
for _, v := range list {
item := &dto.AdminDatasetItem{
Id: v.Id,
Name: v.Name,
Source: v.Source,
ImageCount: v.ImageCount,
LabeledCount: v.LabeledCount,
Status: v.Status,
Cover: v.Cover,
Description: v.Description,
NamePrefix: v.NamePrefix,
CreatedAt: v.CreatedAt,
UpdatedAt: v.UpdatedAt,
}
if t, ok := latest[v.Id]; ok {
item.TrainingId = t.Id
item.TrainingStatus = t.Status
item.TrainingError = t.Error
item.TrainingPublished = published[t.Id]
item.TrainingCurrentEpoch = t.CurrentEpoch
item.TrainingTotalEpochs = t.TotalEpochs
}
items = append(items, item)
}
return &dto.AdminDatasetListRes{Total: total, List: items}, nil
}
// AdminUpdateDataset 更新数据集展示配置(空值字段不覆盖;封面/描述。
// AI 端点/训练机 SSH 走 config.yml 全局配置(localAi / training.ssh))
func (s *datasetService) AdminUpdateDataset(ctx context.Context, req *dto.AdminDatasetUpdateReq) (*dto.AdminDatasetUpdateRes, error) {
existing, err := dao.Dataset.GetById(ctx, req.Id)
if err != nil {
return nil, err
}
if existing == nil {
return nil, gerror.NewCode(common.CodeDatasetNotFound)
}
newName := strings.TrimSpace(req.Name)
if newName != "" && newName != existing.Name {
if err := s.renameDataset(ctx, existing.Id, existing.Name, newName); err != nil {
return nil, err
}
existing.Name = newName
}
// 封面须属于该数据集图片(防伪造文件名指向任意文件)
if req.Cover != "" {
img, err := dao.DatasetImage.GetByFilename(ctx, req.Id, req.Cover)
if err != nil {
return nil, err
}
if img == nil {
return nil, gerror.NewCode(common.CodeImageNotFound)
}
}
prefix := strings.TrimSpace(req.NamePrefix)
if prefix != "" && !datasetNameRe.MatchString(prefix) {
return nil, gerror.New("文件名前缀仅支持字母/数字/下划线/短横线,长度 1-50")
}
if err := dao.Dataset.UpdateConfigs(ctx, req.Id, &entity.Dataset{
Name: newName,
NamePrefix: prefix,
Cover: req.Cover,
Description: req.Description,
}); err != nil {
return nil, err
}
return &dto.AdminDatasetUpdateRes{Id: req.Id}, nil
}
// renameDataset 数据集改名:名称同是磁盘目录名,需同步迁移图片目录与当前生效模型文件。
// 进行中的标注/训练任务持有旧路径,须先等任务结束;改名在 Serial 串行执行,与其他写操作互斥。
func (s *datasetService) renameDataset(ctx context.Context, id int64, oldName, newName string) error {
err := common.Serial().Submit(ctx, func() error {
if !datasetNameRe.MatchString(newName) {
return gerror.New("数据集名称仅支持中文/字母/数字/下划线/短横线,长度 1-50")
}
// 预标注任务进行中(RF-DETR 正在按旧目录读图)
if t, err := dao.LabelTask.GetRunningByDataset(ctx, id); err != nil {
return err
} else if t != nil {
return gerror.NewCode(common.CodeLabelTaskRunning)
}
// 训练进行中(同步目录与产物路径以旧名构建)
if t, err := dao.Training.RunningByDataset(ctx, id); err != nil {
return err
} else if t != nil {
return gerror.New("该数据集有训练任务进行中,无法重命名")
}
// 同名占用检查(UNIQUE 约束兜底)
exists, err := dao.Dataset.GetByName(ctx, newName)
if err != nil {
return err
}
if exists != nil {
return gerror.NewCode(common.CodeDatasetNameDuplicate)
}
// 目录迁移:任一步失败回滚已迁移项(图片目录 / 当前生效模型文件可能不存在,缺省跳过)
oldImgDir := common.DatasetImagesDir(ctx, oldName)
newImgDir := common.DatasetImagesDir(ctx, newName)
oldModel := common.TrainingModelPath(ctx, oldName)
newModel := common.TrainingModelPath(ctx, newName)
movedImgDir, movedModel := false, false
if _, err := os.Stat(oldImgDir); err == nil {
if err := os.Rename(oldImgDir, newImgDir); err != nil {
return gerror.Wrap(err, "迁移图片目录失败")
}
movedImgDir = true
}
if _, err := os.Stat(oldModel); err == nil {
if err := os.Rename(oldModel, newModel); err != nil {
if movedImgDir {
if rbErr := os.Rename(newImgDir, oldImgDir); rbErr != nil {
g.Log().Errorf(ctx, "回滚图片目录失败: %+v", rbErr)
}
}
return gerror.Wrap(err, "迁移模型文件失败")
}
movedModel = true
}
if err := dao.Dataset.UpdateConfigs(ctx, id, &entity.Dataset{Name: newName}); err != nil {
if movedModel {
if rbErr := os.Rename(newModel, oldModel); rbErr != nil {
g.Log().Errorf(ctx, "回滚模型文件失败: %+v", rbErr)
}
}
if movedImgDir {
if rbErr := os.Rename(newImgDir, oldImgDir); rbErr != nil {
g.Log().Errorf(ctx, "回滚图片目录失败: %+v", rbErr)
}
}
return err
}
g.Log().Infof(ctx, "数据集 %d 改名完成: %s -> %s", id, oldName, newName)
return nil
})
return err
}
// AdminUploadCover 上传数据集封面:解码校验图片 → 转 jpg(Quality 92)→ UUID 命名落盘,
// 删除旧封面文件(UUID 互不覆盖,但 DB 指向切换)→ 更新 cover 字段。
func (s *datasetService) AdminUploadCover(ctx context.Context, req *dto.AdminDatasetCoverUploadReq) (*dto.AdminDatasetCoverUploadRes, 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 req.File == nil {
return nil, gerror.New("缺少封面文件")
}
ext := strings.ToLower(filepath.Ext(req.File.Filename))
if !imgExts[ext] {
return nil, gerror.New("封面仅支持 jpg/jpeg/png")
}
if req.File.Size > 2*1024*1024 {
return nil, gerror.New("封面不能超过 2MB")
}
raw, err := req.File.Open()
if err != nil {
return nil, gerror.Wrap(err, "读取封面文件失败")
}
defer func() { _ = raw.Close() }()
img, _, err := image.Decode(raw)
if err != nil {
return nil, gerror.New("封面文件不是有效图片")
}
var buf bytes.Buffer
if err := jpeg.Encode(&buf, img, &jpeg.Options{Quality: 92}); err != nil {
return nil, gerror.Wrap(err, "封面转 jpg 失败")
}
dir := common.DatasetImagesDir(ctx, dataset.Name)
if err := os.MkdirAll(dir, 0o755); err != nil {
return nil, gerror.Wrap(err, "创建图片目录失败")
}
if dataset.Cover != "" {
_ = os.Remove(filepath.Join(dir, filepath.Base(dataset.Cover)))
}
name := common.UuidV4() + ".jpg"
if err := common.WriteFileAtomic(filepath.Join(dir, name), buf.Bytes()); err != nil {
return nil, gerror.Wrap(err, "保存封面失败")
}
if err := dao.Dataset.UpdateConfigs(ctx, dataset.Id, &entity.Dataset{Cover: name}); err != nil {
_ = os.Remove(filepath.Join(dir, name))
return nil, err
}
return &dto.AdminDatasetCoverUploadRes{}, nil
}
// coverNameRe 封面命名规范:UUIDv4 + .jpg(固定命名 cover* 为历史遗留,迁移见 MigrateLegacyCovers
var coverNameRe = regexp.MustCompile(`^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}\.jpg$`)
// isCoverName 校验 DB cover 值是否符合 UUID jpg 规范(防御库内被写入任意文件名)
func isCoverName(name string) bool {
return coverNameRe.MatchString(filepath.Base(name))
}
// CoverFile 封面文件定位(校验归属;controller 直写响应体输出)
func (s *datasetService) CoverFile(ctx context.Context, datasetId int64) (string, error) {
dataset, err := dao.Dataset.GetById(ctx, datasetId)
if err != nil {
return "", err
}
if dataset == nil {
return "", gerror.NewCode(common.CodeDatasetNotFound)
}
if dataset.Cover == "" {
return "", gerror.NewCode(common.CodeImageNotFound)
}
name := filepath.Base(dataset.Cover)
if !isCoverName(name) {
return "", gerror.NewCode(common.CodeImageNotFound)
}
path := filepath.Join(common.DatasetImagesDir(ctx, dataset.Name), name)
if _, err := os.Stat(path); err != nil {
return "", gerror.NewCode(common.CodeImageNotFound)
}
return path, nil
}
// AdminDeleteCover 删除数据集封面:删文件 + 清 cover 字段(列表卡片恢复占位图)
func (s *datasetService) AdminDeleteCover(ctx context.Context, req *dto.AdminDatasetCoverDeleteReq) (*dto.AdminDatasetCoverDeleteRes, 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 dataset.Cover != "" && isCoverName(dataset.Cover) {
_ = os.Remove(filepath.Join(common.DatasetImagesDir(ctx, dataset.Name), filepath.Base(dataset.Cover)))
}
if err := dao.Dataset.ClearCover(ctx, dataset.Id); err != nil {
return nil, err
}
return &dto.AdminDatasetCoverDeleteRes{}, nil
}
// MigrateLegacyCovers 存量封面迁移:cover 非空且不符合 UUID jpg 规范(历史 cover.jpg/cover.png 等)
// → 重命名为 <uuid>.jpg(保留原图字节)并更新 cover 字段。幂等:已符合规范的行跳过。
func (s *datasetService) MigrateLegacyCovers(ctx context.Context) error {
list, err := dao.Dataset.ListAll(ctx)
if err != nil {
return err
}
migrated := 0
for _, d := range list {
if d.Cover == "" || isCoverName(d.Cover) {
continue
}
oldName := filepath.Base(d.Cover)
src := filepath.Join(common.DatasetImagesDir(ctx, d.Name), oldName)
data, err := os.ReadFile(src)
if err != nil {
g.Log().Warningf(ctx, "数据集 %s 封面迁移跳过(文件不存在: %s): %v", d.Name, src, err)
continue
}
newName := common.UuidV4() + ".jpg"
if err := common.WriteFileAtomic(filepath.Join(common.DatasetImagesDir(ctx, d.Name), newName), data); err != nil {
g.Log().Errorf(ctx, "数据集 %s 封面迁移写新文件失败: %+v", d.Name, err)
continue
}
if err := dao.Dataset.UpdateConfigs(ctx, d.Id, &entity.Dataset{Cover: newName}); err != nil {
g.Log().Errorf(ctx, "数据集 %s 封面迁移更新字段失败: %+v", d.Name, err)
continue
}
_ = os.Remove(src)
migrated++
g.Log().Infof(ctx, "数据集 %s 封面迁移: %s → %s", d.Name, oldName, newName)
}
if migrated > 0 {
g.Log().Infof(ctx, "封面存量迁移完成: 共 %d 个数据集", migrated)
}
return nil
}
// AdminCreateDataset 新建数据集:名称唯一(UNIQUE 兜底)+ 创建图片目录
func (s *datasetService) AdminCreateDataset(ctx context.Context, req *dto.AdminDatasetCreateReq) (*dto.AdminDatasetCreateRes, error) {
now := gtime.Now()
var id int64
err := common.Serial().Submit(ctx, func() error {
exists, err := dao.Dataset.GetByName(ctx, req.Name)
if err != nil {
return err
}
if exists != nil {
return gerror.NewCode(common.CodeDatasetNameDuplicate)
}
prefix := strings.TrimSpace(req.NamePrefix)
if prefix != "" && !datasetNameRe.MatchString(prefix) {
return gerror.New("文件名前缀仅支持字母/数字/下划线/短横线,长度 1-50")
}
id, err = dao.Dataset.Insert(ctx, &entity.Dataset{
Name: req.Name,
Source: req.Source,
NamePrefix: prefix,
Status: consts.DatasetStatusBuilding,
CreatedAt: now,
UpdatedAt: now,
})
return err
})
if err != nil {
return nil, err
}
if err := os.MkdirAll(common.DatasetImagesDir(ctx, req.Name), 0o755); err != nil {
return nil, gerror.Wrap(err, "创建数据集目录失败")
}
return &dto.AdminDatasetCreateRes{Id: id}, nil
}
// AdminDeleteDataset 删除数据集:有 running 标注任务 / 该数据集训练进行中 / 已发布模型版本时拒绝
// (训练产物与模型为付费资产,需先删除模型版本再删数据集)。
// 删除 = 删图片/模型目录 + 删记录(标注随图片行删除,Serial 单写者串行)。
func (s *datasetService) AdminDeleteDataset(ctx context.Context, req *dto.AdminDatasetDeleteReq) (*dto.AdminDatasetDeleteRes, error) {
var name string
err := common.Serial().Submit(ctx, func() error {
d, err := dao.Dataset.GetById(ctx, req.Id)
if err != nil {
return err
}
if d == nil {
return gerror.NewCode(common.CodeDatasetNotFound)
}
name = d.Name
// 预标注任务进行中(RF-DETR 正在扫该数据集图片)
if t, err := dao.LabelTask.GetRunningByDataset(ctx, d.Id); err != nil {
return err
} else if t != nil {
return gerror.NewCode(common.CodeLabelTaskRunning)
}
// 该数据集训练进行中(并发度 1,防文件被删训练中断)
if t, err := dao.Training.RunningByDataset(ctx, d.Id); err != nil {
return err
} else if t != nil {
return gerror.New("该数据集有训练任务进行中,无法删除")
}
// 模型版本记录随数据集级联删除(管理端无模型管理界面,2026-08-26 决策;
// 若需保留已下发模型,删除数据集前先确认客户端不再需要)
if err := dao.ModelVersion.DeleteByDataset(ctx, d.Id); err != nil {
return err
}
if err := dao.DatasetImage.DeleteByDataset(ctx, d.Id); err != nil {
return err
}
return dao.Dataset.DeleteById(ctx, d.Id)
})
if err != nil {
return nil, err
}
// 文件清理(图片目录 + 当前生效模型文件;删除失败仅记日志,记录已删)
paths := []string{common.DatasetImagesDir(ctx, name), common.TrainingModelPath(ctx, name)}
for _, p := range paths {
if err := os.RemoveAll(p); err != nil {
g.Log().Errorf(ctx, "删除数据集 %s 目录失败: %+v", p, err)
}
}
return &dto.AdminDatasetDeleteRes{}, nil
}
// AdminUploadImages 上传图片(多文件):重名跳过(文件名唯一约束),逐张落盘 + 批量入库。
func (s *datasetService) AdminUploadImages(ctx context.Context, req *dto.AdminDatasetUploadReq) (*dto.AdminDatasetUploadRes, error) {
if len(req.Files) == 0 {
return nil, gerror.New("请选择图片文件")
}
dataset, err := dao.Dataset.GetById(ctx, req.DatasetId)
if err != nil {
return nil, err
}
if dataset == nil {
return nil, gerror.NewCode(common.CodeDatasetNotFound)
}
dir := common.DatasetImagesDir(ctx, dataset.Name)
if err := os.MkdirAll(dir, 0o755); err != nil {
return nil, gerror.Wrap(err, "创建图片目录失败")
}
res := &dto.AdminDatasetUploadRes{}
now := gtime.Now()
saved := make([]string, 0, len(req.Files)) // 本次已落盘文件名(入库失败时清理)
for _, f := range req.Files {
name := filepath.Base(f.Filename)
if !imgExts[strings.ToLower(filepath.Ext(name))] {
continue
}
exists, err := dao.DatasetImage.GetByFilename(ctx, dataset.Id, name)
if err != nil {
return nil, err
}
if exists != nil {
res.Skipped++
continue
}
if err := saveImageFile(f, dir); err != nil {
return nil, gerror.Wrapf(err, "图片 %s 保存失败", name)
}
saved = append(saved, name)
}
if len(saved) == 0 {
return res, nil
}
// 标注强语义:未配置标注服务时不允许产生无标注图(文件已落盘,失败则清掉)
if common.LocalAiClient(ctx) == nil {
for _, name := range saved {
_ = os.Remove(filepath.Join(dir, name))
}
return nil, gerror.NewCode(common.CodeLocalAiNotConfigured)
}
addedIds := make([]int64, 0, len(saved))
err = common.Serial().Submit(ctx, func() error {
for _, name := range saved {
// 重查重(并发上传兜底)+ 入库
exists, err := dao.DatasetImage.GetByFilename(ctx, dataset.Id, name)
if err != nil {
return err
}
if exists != nil {
continue
}
id, err := dao.DatasetImage.Insert(ctx, &entity.DatasetImage{
DatasetId: dataset.Id,
Filename: name,
Source: "manual",
CreatedAt: now,
})
if err != nil {
return err
}
addedIds = append(addedIds, id)
res.Added++
}
return dao.Dataset.UpdateCounters(ctx, dataset.Id, int64(res.Added), -1, "")
})
if err != nil {
// 入库失败:清掉已落盘文件,保证「记录存在 ⟺ 文件存在」
for _, name := range saved {
_ = os.Remove(filepath.Join(dir, name))
}
return nil, err
}
// 自动触发标注:忙(已有 running 任务)不报错,由任务完成后的自动补标轮兜底;其他失败回滚本次入库
newImages, err := dao.DatasetImage.GetByIds(ctx, addedIds)
if err != nil {
return nil, err
}
if err := LabelTask.AutoLabel(ctx, dataset, newImages); err != nil {
// 触发失败回滚本次入库(manual 上传文件非付费资产,可删)
rollbackIds := make([]int64, 0, len(newImages))
for _, img := range newImages {
rollbackIds = append(rollbackIds, img.Id)
if rErr := os.Remove(filepath.Join(dir, img.Filename)); rErr != nil {
g.Log().Warningf(ctx, "回滚删除图片文件失败: %s: %v", img.Filename, rErr)
}
}
if rErr := common.Serial().Submit(ctx, func() error {
if dErr := dao.DatasetImage.DeleteByIds(ctx, rollbackIds); dErr != nil {
return dErr
}
return dao.Dataset.UpdateCounters(ctx, dataset.Id, -int64(len(rollbackIds)), 0, "")
}); rErr != nil {
g.Log().Errorf(ctx, "自动标注触发失败后的入库回滚失败: %v", rErr)
}
return nil, err
}
return res, nil
}
// saveImageFile 上传文件保存(ghttp.UploadFile.Save 已按原始文件名落盘)
func saveImageFile(f *ghttp.UploadFile, dir string) error {
saved, err := f.Save(dir)
if err != nil {
return err
}
if saved != filepath.Base(f.Filename) {
_ = os.Remove(filepath.Join(dir, saved))
}
return nil
}
// nextImageIndex 生成图顺序编号起点:数据集内 <前缀>_<数字>.jpg 的已有最大序号 + 1
func (s *datasetService) nextImageIndex(ctx context.Context, datasetId int64, prefix string) (int, error) {
list, err := dao.DatasetImage.ListByDataset(ctx, datasetId)
if err != nil {
return 0, err
}
re := regexp.MustCompile(`^` + regexp.QuoteMeta(prefix) + `_(\d+)\.jpg$`)
max := 0
for _, img := range list {
if m := re.FindStringSubmatch(img.Filename); m != nil {
if v, err := strconv.Atoi(m[1]); err == nil && v > max {
max = v
}
}
}
return max + 1, nil
}
// validateGeneratedImage 生成图距离/数量校验:RF-DETR 全图检测目标框(与自动标注同口径去重),
// 数量精确匹配 animalCount;距离用物理公式 d = 物种体高 × focalPx ÷ 目标像素高
// focalPx 由假设视场角 imageGen.assumedVfovDeg 推出;体高按物种查 imageGen.speciesHeights),
// 须落在表单 [distanceMin, distanceMax] 内(>0 才启用对应项)
func (s *datasetService) validateGeneratedImage(ctx context.Context, data []byte, req *dto.AdminDatasetGenerateReq, species string) (bool, string, error) {
client := common.LocalAiClient(ctx)
if client == nil {
return false, "", gerror.NewCode(common.CodeLocalAiNotConfigured)
}
// 校验与预标注用途不同:预标注阈值 0.08 宁多勿漏,会把玉米叶/土块算成目标(实测 30-40m 场景 6-18 框噪声);
// 校验用高置信度阈值(实测 0.3 时真鸟全留、噪声清零,0.4 起连远处真鸟都丢)
client.Threshold = g.Cfg().MustGet(ctx, "imageGen.validateThreshold", 0.3).Float64()
imgW, imgH, err := parseGenSize(req.Size)
if err != nil {
return false, "", err
}
dets, err := client.Detect(ctx, data, "image/jpeg", imgW, imgH)
if err != nil {
return false, "", gerror.Wrap(err, "生成图检测失败")
}
// 与自动标注同口径去重(同目标多框只计一只),保证「校验数量 == 标注框数」
cands := make([]*dto.AdminLabelBox, 0, len(dets))
for _, d := range dets {
cands = append(cands, &dto.AdminLabelBox{
Cx: (d.X + d.Width/2) / float64(imgW),
Cy: (d.Y + d.Height/2) / float64(imgH),
W: d.Width / float64(imgW),
H: d.Height / float64(imgH),
Confidence: d.Confidence,
})
}
deduped := suppressOverlap(cands, client.OverlapThreshold)
if len(deduped) == 0 {
return false, "未检测到目标动物", nil
}
// 数量为精确校验:表单指定几只,画面里就必须恰好几只(提示词中的数量应与表单一致)
if req.AnimalCount > 0 && len(deduped) != req.AnimalCount {
return false, fmt.Sprintf("目标数量 %d 与指定数量 %d 不符", len(deduped), req.AnimalCount), nil
}
maxRatio := 0.0
for _, b := range deduped {
if b.H > maxRatio {
maxRatio = b.H
}
}
// 物理公式:d = 体高 × focalPx ÷ 像素高;focalPx = 图高/2 ÷ tan(vfov/2)
heightM := g.Cfg().MustGet(ctx, "imageGen.speciesHeights.default", 0.35).Float64()
if h := g.Cfg().MustGet(ctx, "imageGen.speciesHeights."+species).Float64(); h > 0 {
heightM = h
}
vfov := g.Cfg().MustGet(ctx, "imageGen.assumedVfovDeg", 52.0).Float64() * math.Pi / 180
focalPx := float64(imgH) / 2 / math.Tan(vfov/2)
pxHeight := maxRatio * float64(imgH)
estDist := heightM * focalPx / pxHeight
if req.DistanceMin > 0 && estDist < float64(req.DistanceMin) {
return false, fmt.Sprintf("估算距离约 %.0f 米低于下限 %d 米(目标过近)", estDist, req.DistanceMin), nil
}
if req.DistanceMax > 0 && estDist > float64(req.DistanceMax) {
return false, fmt.Sprintf("估算距离约 %.0f 米超出上限 %d 米(目标过远)", estDist, req.DistanceMax), nil
}
return true, fmt.Sprintf("数量 %d,最大目标估算距离约 %.0f 米", len(deduped), estDist), nil
}
// parseGenSize 解析 "704x1248" 形式尺寸
func parseGenSize(size string) (int, int, error) {
parts := strings.SplitN(size, "x", 2)
if len(parts) != 2 {
return 0, 0, gerror.Newf("非法图片尺寸 %q", size)
}
w, errW := strconv.Atoi(parts[0])
h, errH := strconv.Atoi(parts[1])
if errW != nil || errH != nil || w <= 0 || h <= 0 {
return 0, 0, gerror.Newf("非法图片尺寸 %q", size)
}
return w, h, nil
}
// countWordCN 数量词(与表单 animalCount 精确对应,防提示词文案与校验数量不符)
var countWordCN = [...]string{"一只", "两只", "三只", "四只", "五只", "六只", "七只", "八只", "九只", "十只"}
// AdminGenOptions 生成表单选项:该数据集的物种池 + 场景池(config 维护,前端下拉用)
func (s *datasetService) AdminGenOptions(ctx context.Context, req *dto.AdminDatasetGenOptionsReq) (*dto.AdminDatasetGenOptionsRes, error) {
dataset, err := dao.Dataset.GetById(ctx, req.DatasetId)
if err != nil {
return nil, err
}
if dataset == nil {
return nil, gerror.NewCode(common.CodeDatasetNotFound)
}
return &dto.AdminDatasetGenOptionsRes{
Species: g.Cfg().MustGet(ctx, "imageGen.speciesByDataset."+dataset.Name).Strings(),
Scenes: g.Cfg().MustGet(ctx, "imageGen.scenes").Strings(),
}, nil
}
// buildPromptFromTemplate 按 config 通用模板组装提示词:物种/场景优先用表单选择
// (物种只写名字不写羽毛细节——模型对正确名称自带外观先验;场景未选则池内随机),
// 物种未填按数据集名取物种池随机(池也未配置则用数据集名本身);动作/光线池内随机。
func (s *datasetService) buildPromptFromTemplate(ctx context.Context, datasetName, species string, count int) (string, string, error) {
tpl := g.Cfg().MustGet(ctx, "imageGen.promptTemplate").String()
if tpl == "" {
return "", "", gerror.New("未配置通用提示词模板(imageGen.promptTemplate),请手填提示词")
}
if species == "" {
species = datasetName
if pool := g.Cfg().MustGet(ctx, "imageGen.speciesByDataset."+datasetName).Strings(); len(pool) > 0 {
species = pool[rand.Intn(len(pool))]
}
}
scene := pickCfg(ctx, "imageGen.scenes")
action := pickCfg(ctx, "imageGen.actions")
light := pickCfg(ctx, "imageGen.lights")
dist := g.Cfg().MustGet(ctx, "imageGen.distanceWord", "远处").String()
countWord := fmt.Sprintf("%d只", count)
if count >= 1 && count <= len(countWordCN) {
countWord = countWordCN[count-1]
}
out := tpl
for k, v := range map[string]string{
"{scene}": scene, "{species}": species, "{count}": countWord,
"{action}": action, "{light}": light, "{distanceWord}": dist,
} {
out = strings.ReplaceAll(out, k, v)
}
g.Log().Infof(ctx, "模板组装提示词: %s", out)
return out, species, nil
}
// pickCfg 取字符串数组配置项随机一项(空池返回空串)
func pickCfg(ctx context.Context, key string) string {
pool := g.Cfg().MustGet(ctx, key).Strings()
if len(pool) == 0 {
return ""
}
return pool[rand.Intn(len(pool))]
}
// AdminGenerateImages AI 生成图片:同步逐张生成(每张超时 imageGen.timeoutSeconds),逐张落盘 + 入库;
// 中途失败返回错误并附已成功张数(已生成的图为付费资产,保留不删除)。
func (s *datasetService) AdminGenerateImages(ctx context.Context, req *dto.AdminDatasetGenerateReq) (*dto.AdminDatasetGenerateRes, error) {
provider := common.ImageGen(ctx)
if provider == nil {
return nil, gerror.NewCode(common.CodeImageGenNotConfigured)
}
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)
}
dir := common.DatasetImagesDir(ctx, dataset.Name)
if err := os.MkdirAll(dir, 0o755); err != nil {
return nil, gerror.Wrap(err, "创建图片目录失败")
}
now := gtime.Now()
generated := 0
addedIds := make([]int64, 0, req.Count)
// 每张生成超时走配置(dashscope 默认 120slocalai 训练机约 5min/张,建议 600)
genTimeout := time.Duration(g.Cfg().MustGet(ctx, "imageGen.timeoutSeconds", 120).Int()) * time.Second
handWritten := strings.TrimSpace(req.Prompt) != ""
genOnce := func(prompt string) ([]byte, error) {
genCtx, cancel := context.WithTimeout(ctx, genTimeout)
defer cancel()
return provider.Generate(genCtx, prompt, req.Size)
}
// 命名:数据集属性 name_prefix(创建/编辑表单维护)非空时按 <前缀>_<两位序号>.jpg 顺序编号
// (延续数据集内已有最大序号,参考野鸡目录 pheasant_NN.jpg 惯例);无前缀按时间戳命名。
// 注意:gtime.Time.Format 是 PHP 风格格式符,Go 数字布局须用 .Time.Format,否则原样透传
prefix := dataset.NamePrefix
nextIdx := 1
if prefix != "" {
if !datasetNameRe.MatchString(prefix) {
return nil, gerror.New("文件名前缀仅支持字母/数字/下划线/短横线")
}
nextIdx, err = s.nextImageIndex(ctx, dataset.Id, prefix)
if err != nil {
return nil, err
}
}
for i := 0; i < req.Count; i++ {
// 部分成功也要标注:报错返回前对已入库图片触发自动标注(外部重试驱动依赖此语义攒量)
failWith := func(err error) (*dto.AdminDatasetGenerateRes, error) {
if generated > 0 && len(addedIds) > 0 {
if newImages, lErr := dao.DatasetImage.GetByIds(ctx, addedIds); lErr == nil {
if aErr := LabelTask.AutoLabel(ctx, dataset, newImages); aErr != nil {
g.Log().Errorf(ctx, "部分成功后自动标注触发失败: %+v", aErr)
}
}
}
return &dto.AdminDatasetGenerateRes{Generated: generated}, err
}
// 提示词:手填优先;未手填走 config 通用模板逐张组装——物种/场景用表单选择,
// 未填按数据集名取物种池随机、场景池随机;数量词与表单 animalCount 一致;动作/光线随机
prompt := strings.TrimSpace(req.Prompt)
species := strings.TrimSpace(req.Species)
if !handWritten {
var tErr error
if prompt, species, tErr = s.buildPromptFromTemplate(ctx, dataset.Name, species, req.AnimalCount); tErr != nil {
return nil, tErr
}
}
data, genErr := genOnce(prompt)
if genErr != nil {
// 付费资产保留原则:已生成的不删除
return failWith(gerror.NewCode(common.CodeImageGenFailed, fmt.Sprintf("第 %d 张生成失败(已生成 %d 张): %v", i+1, generated, genErr)))
}
// 逐张校验(表单填写距离范围/目标数量后启用):不合格丢弃并重新生成,连续不合格报错返回
needValidate := req.DistanceMax > 0 || req.DistanceMin > 0 || req.AnimalCount > 0
if needValidate {
passed := false
reason := ""
for attempt := 1; attempt <= consts.GenValidateMaxAttempts; attempt++ {
ok, r, vErr := s.validateGeneratedImage(ctx, data, req, species)
if vErr != nil {
return failWith(gerror.Wrapf(vErr, "第 %d 张校验失败(已生成 %d 张)", i+1, generated))
}
if ok {
passed = true
break
}
reason = r
g.Log().Infof(ctx, "生成图校验不合格(%s),第 %d 张第 %d 次重新生成", r, i+1, attempt)
if data, genErr = genOnce(prompt); genErr != nil {
return failWith(gerror.NewCode(common.CodeImageGenFailed, fmt.Sprintf("第 %d 张重新生成失败(已生成 %d 张): %v", i+1, generated, genErr)))
}
}
if !passed {
return failWith(gerror.NewCode(common.CodeImageGenFailed, fmt.Sprintf("第 %d 张连续 %d 次校验不合格(%s),请调整提示词或范围后重试(已生成 %d 张)", i+1, consts.GenValidateMaxAttempts, reason, generated)))
}
}
filename := ""
if prefix != "" {
filename = fmt.Sprintf("%s_%02d.jpg", prefix, nextIdx)
nextIdx++
} else {
filename = fmt.Sprintf("gen_%s_%d.jpg", now.Time.Format("20060102150405"), i)
}
if err := common.WriteFileAtomic(filepath.Join(dir, filename), data); err != nil {
return failWith(gerror.Wrapf(err, "第 %d 张保存失败(已生成 %d 张)", i+1, generated))
}
var imgId int64
insErr := common.Serial().Submit(ctx, func() error {
id, err := dao.DatasetImage.Insert(ctx, &entity.DatasetImage{
DatasetId: dataset.Id,
Filename: filename,
Source: "ai",
Prompt: prompt,
AnimalCount: req.AnimalCount,
CreatedAt: now,
})
imgId = id
return err
})
if insErr != nil {
return failWith(gerror.Wrap(insErr, "生成图片入库失败"))
}
addedIds = append(addedIds, imgId)
generated++
}
if err := common.Serial().Submit(ctx, func() error {
return dao.Dataset.UpdateCounters(ctx, dataset.Id, int64(generated), -1, "")
}); err != nil {
return nil, err
}
// 自动触发标注:忙(已有 running 任务)不报错,由任务完成后的自动补标轮兜底;
// 其他失败报错但保留已生成图(付费资产,不可删)
newImages, err := dao.DatasetImage.GetByIds(ctx, addedIds)
if err != nil {
return nil, err
}
if err := LabelTask.AutoLabel(ctx, dataset, newImages); err != nil {
return &dto.AdminDatasetGenerateRes{Generated: generated},
gerror.Wrap(err, "图片已生成入库,但自动标注触发失败")
}
return &dto.AdminDatasetGenerateRes{Generated: generated}, nil
}
// AdminListImages 数据集图片列表(创建时间正序,标注工作台/网格预览)
func (s *datasetService) AdminListImages(ctx context.Context, req *dto.AdminDatasetImagesReq) (*dto.AdminDatasetImagesRes, error) {
dataset, err := dao.Dataset.GetById(ctx, req.DatasetId)
if err != nil {
return nil, err
}
if dataset == nil {
return nil, gerror.NewCode(common.CodeDatasetNotFound)
}
list, err := dao.DatasetImage.ListByDataset(ctx, dataset.Id)
if err != nil {
return nil, err
}
items := make([]*dto.AdminImageItem, 0, len(list))
for _, v := range list {
items = append(items, &dto.AdminImageItem{
Id: v.Id,
Filename: v.Filename,
Source: v.Source,
Prompt: v.Prompt,
Url: datasetImageUrl(ctx, dataset.Id, v.Filename),
CreatedAt: v.CreatedAt,
})
}
return &dto.AdminDatasetImagesRes{List: items}, nil
}
// AdminDeleteImages 删除图片:删图片文件 + 记录 + 更新计数
// (付费资产,前端带确认文案;标注/候选框 JSON 随图片行删除)
func (s *datasetService) AdminDeleteImages(ctx context.Context, req *dto.AdminDatasetImagesDeleteReq) (*dto.AdminDatasetImagesDeleteRes, error) {
dataset, err := dao.Dataset.GetById(ctx, req.DatasetId)
if err != nil {
return nil, err
}
if dataset == nil {
return nil, gerror.NewCode(common.CodeDatasetNotFound)
}
images, err := dao.DatasetImage.GetByIds(ctx, req.Ids)
if err != nil {
return nil, err
}
var del []*entity.DatasetImage
for _, img := range images {
// 只删本数据集内的记录(防跨数据集误删)
if img.DatasetId == dataset.Id {
del = append(del, img)
}
}
if len(del) == 0 {
return &dto.AdminDatasetImagesDeleteRes{}, nil
}
// 文件清理(删除失败仅记日志,记录照删):图片文件(标注/候选随行删除)
imgDir := common.DatasetImagesDir(ctx, dataset.Name)
for _, img := range del {
_ = os.Remove(filepath.Join(imgDir, img.Filename))
}
// Serial 内:删行 → 重算已标注数 → 同步计数与状态
err = common.Serial().Submit(ctx, func() error {
ids := make([]int64, 0, len(del))
for _, img := range del {
ids = append(ids, img.Id)
}
if err := dao.DatasetImage.DeleteByIds(ctx, ids); 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, -int64(len(del)), labeled, status)
})
if err != nil {
return nil, err
}
return &dto.AdminDatasetImagesDeleteRes{}, nil
}
// datasetImageUrl 管理端图片访问地址(controller 直写响应体,见 admin.go Image 方法)
func datasetImageUrl(ctx context.Context, datasetId int64, filename string) string {
return fmt.Sprintf("/api/v1/admin/datasets/image?%s",
url.Values{"datasetId": {fmt.Sprintf("%d", datasetId)}, "filename": {filename}}.Encode())
}
// ImageFile 图片文件定位(校验归属后返回绝对路径;controller 直写响应体)
func (s *datasetService) ImageFile(ctx context.Context, datasetId int64, filename string) (string, error) {
if filename != filepath.Base(filename) || !imgExts[strings.ToLower(filepath.Ext(filename))] {
return "", gerror.NewCode(common.CodeImageNotFound)
}
dataset, err := dao.Dataset.GetById(ctx, datasetId)
if err != nil {
return "", err
}
if dataset == nil {
return "", gerror.NewCode(common.CodeDatasetNotFound)
}
img, err := dao.DatasetImage.GetByFilename(ctx, dataset.Id, filename)
if err != nil {
return "", err
}
if img == nil {
return "", gerror.NewCode(common.CodeImageNotFound)
}
path := filepath.Join(common.DatasetImagesDir(ctx, dataset.Name), filename)
if _, err := os.Stat(path); err != nil {
return "", gerror.NewCode(common.CodeImageNotFound)
}
return path, nil
}