1797 lines
70 KiB
Go
1797 lines
70 KiB
Go
package service
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"image"
|
||
"image/jpeg"
|
||
"io"
|
||
"math"
|
||
"math/bits"
|
||
"math/rand"
|
||
"net/url"
|
||
"os"
|
||
"path/filepath"
|
||
"regexp"
|
||
"sort"
|
||
"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"
|
||
"golang.org/x/image/draw"
|
||
|
||
"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)
|
||
}
|
||
// 图片数/已标注数实时统计（冗余计数列 image_count/labeled_count 曾因生成任务部分失败
|
||
// 未回写等路径漂移为负数，展示一律以 dataset_image 实际统计为准）
|
||
stats, err := dao.DatasetImage.CountStatByDatasets(ctx, ids)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
latest, err := dao.Training.LatestByDatasets(ctx, ids)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
// 已发布过的训练不再返回发布按钮（model_version.training_id 反查）
|
||
trainingIds := make([]int64, 0, len(latest))
|
||
for _, ts := range latest {
|
||
for _, t := range ts {
|
||
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 {
|
||
st := stats[v.Id] // 无图片行时为零值（计数 0），不落冗余列
|
||
item := &dto.AdminDatasetItem{
|
||
Id: v.Id,
|
||
Name: v.Name,
|
||
Source: v.Source,
|
||
ImageCount: st.Count,
|
||
LabeledCount: st.Labeled,
|
||
Status: v.Status,
|
||
Cover: v.Cover,
|
||
Description: v.Description,
|
||
NamePrefix: v.NamePrefix,
|
||
GenSpecies: v.GenSpecies,
|
||
GenTone: v.GenTone,
|
||
GenHeights: v.GenHeights,
|
||
GenScenes: v.GenScenes,
|
||
GenActions: v.GenActions,
|
||
GenOcclusions: v.GenOcclusions,
|
||
GenClasses: v.GenClasses,
|
||
SortOrder: v.SortOrder,
|
||
Trains: []*dto.AdminDatasetTrainBrief{},
|
||
CreatedAt: v.CreatedAt,
|
||
UpdatedAt: v.UpdatedAt,
|
||
}
|
||
for _, t := range latest[v.Id] {
|
||
item.Trains = append(item.Trains, &dto.AdminDatasetTrainBrief{
|
||
TrainingId: t.Id,
|
||
Variant: t.Variant,
|
||
Status: t.Status,
|
||
Error: t.Error,
|
||
Published: published[t.Id],
|
||
CurrentEpoch: t.CurrentEpoch,
|
||
TotalEpochs: 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)
|
||
prefix := strings.TrimSpace(req.NamePrefix)
|
||
if prefix != "" && !datasetNameRe.MatchString(prefix) {
|
||
return nil, gerror.New("文件名前缀仅支持字母/数字/下划线/短横线，长度 1-50")
|
||
}
|
||
if newName != "" && newName != existing.Name {
|
||
// 改名（可同时改前缀）：图片目录 + 模型文件随命名迁移（模型文件名 = 前缀，空回退数据集名）
|
||
if err := s.renameDataset(ctx, existing.Id, existing.Name, newName, existing.NamePrefix, prefix); err != nil {
|
||
return nil, err
|
||
}
|
||
existing.Name = newName
|
||
} else if prefix != existing.NamePrefix {
|
||
// 仅前缀变更：模型文件名随前缀迁移（改名场景已由 renameDataset 处理）
|
||
migrateModelFile(ctx, existing.Name, existing.NamePrefix, existing.Name, prefix)
|
||
}
|
||
// 封面须属于该数据集图片（防伪造文件名指向任意文件）
|
||
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)
|
||
}
|
||
}
|
||
if err := dao.Dataset.UpdateConfigs(ctx, req.Id, &entity.Dataset{
|
||
Name: newName,
|
||
NamePrefix: prefix,
|
||
Cover: req.Cover,
|
||
}); err != nil {
|
||
return nil, err
|
||
}
|
||
// 描述允许清空（UpdateConfigs 空值不覆盖），单独更新
|
||
if err := dao.Dataset.UpdateDescription(ctx, req.Id, strings.TrimSpace(req.Description)); err != nil {
|
||
return nil, err
|
||
}
|
||
// 序号独立更新（指针区分未传与显式 0；0 时按创建时间倒序排）
|
||
if req.SortOrder != nil {
|
||
if err := dao.Dataset.UpdateSortOrder(ctx, req.Id, *req.SortOrder); err != nil {
|
||
return nil, err
|
||
}
|
||
}
|
||
// gen_* 7 列不在此维护：创建数据集时 VLM 生成，重新生成走 gen-pools 接口
|
||
return &dto.AdminDatasetUpdateRes{Id: req.Id}, nil
|
||
}
|
||
|
||
// renameDataset 数据集改名：名称同是磁盘目录名，需同步迁移图片目录；
|
||
// 模型文件按文件名前缀命名（空回退数据集名），随改名/改前缀迁移（失败仅记日志不阻断——旧文件保留）。
|
||
// 进行中的标注/训练任务持有旧路径，须先等任务结束；改名在 Serial 串行执行，与其他写操作互斥。
|
||
func (s *datasetService) renameDataset(ctx context.Context, id int64, oldName, newName, oldPrefix, newPrefix 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)
|
||
movedImgDir := false
|
||
if _, err := os.Stat(oldImgDir); err == nil {
|
||
if err := os.Rename(oldImgDir, newImgDir); err != nil {
|
||
return gerror.Wrap(err, "迁移图片目录失败")
|
||
}
|
||
movedImgDir = true
|
||
}
|
||
// 模型文件随命名迁移（prefix 空回退数据集名；改名不改前缀时路径不变，内部跳过）
|
||
migrateModelFile(ctx, oldName, oldPrefix, newName, newPrefix)
|
||
if err := dao.Dataset.UpdateConfigs(ctx, id, &entity.Dataset{Name: newName}); err != nil {
|
||
// DB 失败回滚已迁移项（模型文件反向迁移失败仅记日志，旧文件可能留新名）
|
||
migrateModelFile(ctx, newName, newPrefix, oldName, oldPrefix)
|
||
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
|
||
}
|
||
|
||
// migrateModelFile 模型文件随命名变更迁移（基名 = 文件名前缀，空回退数据集名）：
|
||
// 双档位（2026-09-03）s/n 两文件一并迁移（n 档文件名带 _n 后缀）；
|
||
// 旧文件不存在或新旧路径相同（前缀未变且非空）直接跳过；失败仅记日志不阻断业务——
|
||
// 模型文件为付费训练产物，保留旧名总比删除好（下次训练直写新名覆盖）。
|
||
func migrateModelFile(ctx context.Context, oldName, oldPrefix, newName, newPrefix string) {
|
||
migrateModelFileVariant(ctx, oldName, oldPrefix, newName, newPrefix, "")
|
||
migrateModelFileVariant(ctx, oldName, oldPrefix, newName, newPrefix, consts.TrainingVariantNFileSuffix)
|
||
}
|
||
|
||
// migrateModelFileVariant 迁移单档位文件（variantSuffix 空=s 档，_n=n 档）
|
||
func migrateModelFileVariant(ctx context.Context, oldName, oldPrefix, newName, newPrefix, variantSuffix string) {
|
||
oldPath := common.TrainingModelPath(ctx, modelFileName(oldName, oldPrefix)+variantSuffix)
|
||
newPath := common.TrainingModelPath(ctx, modelFileName(newName, newPrefix)+variantSuffix)
|
||
if oldPath == newPath {
|
||
return
|
||
}
|
||
if _, err := os.Stat(oldPath); err != nil {
|
||
return
|
||
}
|
||
if err := os.Rename(oldPath, newPath); err != nil {
|
||
g.Log().Errorf(ctx, "迁移模型文件 %s → %s 失败: %+v", oldPath, newPath, 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 > 10*1024*1024 {
|
||
return nil, gerror.New("封面不能超过 10MB")
|
||
}
|
||
raw, err := req.File.Open()
|
||
if err != nil {
|
||
return nil, gerror.Wrap(err, "读取封面文件失败")
|
||
}
|
||
defer func() { _ = raw.Close() }()
|
||
data, err := io.ReadAll(raw)
|
||
if err != nil {
|
||
return nil, gerror.Wrap(err, "读取封面文件失败")
|
||
}
|
||
// 统一压缩为 1248x704（任意尺寸源图 → 目标尺寸 jpg）
|
||
data, err = resizeCover(data)
|
||
if err != nil {
|
||
return nil, gerror.Wrap(err, "封面文件不是有效图片")
|
||
}
|
||
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), data); 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 直写响应体输出）。
|
||
// datasetId>0 从库取封面名；datasetId=0 时按 name+filename 直读（新建对话框预生成封面回显，
|
||
// 数据集未创建无库记录；filename 须符合 UUID jpg 规范防路径穿越）。
|
||
func (s *datasetService) CoverFile(ctx context.Context, datasetId int64, name, filename string) (string, error) {
|
||
if datasetId > 0 {
|
||
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 = dataset.Name
|
||
filename = dataset.Cover
|
||
}
|
||
name = filepath.Base(name)
|
||
if !datasetNameRe.MatchString(name) {
|
||
return "", gerror.NewCode(common.CodeImageNotFound)
|
||
}
|
||
filename = filepath.Base(filename)
|
||
if !isCoverName(filename) {
|
||
return "", gerror.NewCode(common.CodeImageNotFound)
|
||
}
|
||
path := filepath.Join(common.DatasetImagesDir(ctx, name), filename)
|
||
if _, err := os.Stat(path); err != nil {
|
||
return "", gerror.NewCode(common.CodeImageNotFound)
|
||
}
|
||
return path, nil
|
||
}
|
||
|
||
// AppCoverFile 公开封面定位（App 模型清单缩略图）：按训练文件名前缀查数据集，
|
||
// 复用 CoverFile 的归属校验（数据集不存在/无封面/文件缺失 → 404 错误码）。
|
||
func (s *datasetService) AppCoverFile(ctx context.Context, namePrefix string) (string, error) {
|
||
d, err := dao.Dataset.GetByNamePrefix(ctx, strings.TrimSpace(namePrefix))
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
if d == nil {
|
||
return "", gerror.NewCode(common.CodeDatasetNotFound)
|
||
}
|
||
return s.CoverFile(ctx, d.Id, "", "")
|
||
}
|
||
|
||
// 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
|
||
}
|
||
|
||
// AdminGenCover 生成数据集封面（对话框「AI 生成封面」按钮入口）：
|
||
// datasetId>0 走完整链路（16:9、1 雄 1 雌，物种取 gen_species 空回退数据集名，写库+删旧封面）；
|
||
// datasetId=0 新建预生成（物种=表单数据集名，仅落盘，创建请求带 cover 回传写库）；
|
||
// 成功返回新 cover 文件名供前端回显；失败报错不覆盖旧封面。
|
||
func (s *datasetService) AdminGenCover(ctx context.Context, req *dto.AdminGenCoverReq) (*dto.AdminGenCoverRes, error) {
|
||
datasetName := strings.TrimSpace(req.Name)
|
||
species := ""
|
||
if req.DatasetId > 0 {
|
||
dataset, err := dao.Dataset.GetById(ctx, req.DatasetId)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if dataset == nil {
|
||
return nil, gerror.NewCode(common.CodeDatasetNotFound)
|
||
}
|
||
datasetName = dataset.Name
|
||
species = dataset.GenSpecies
|
||
}
|
||
if !datasetNameRe.MatchString(datasetName) {
|
||
return nil, gerror.New("缺少有效的数据集名称（新建时请先填写名称再生成封面）")
|
||
}
|
||
cover, err := s.genCoverWithImageGen(ctx, req.DatasetId, datasetName, species)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return &dto.AdminGenCoverRes{Cover: cover}, 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
|
||
}
|
||
|
||
// CompressExistingCovers 存量封面尺寸统一迁移：cover 非空且尺寸不是 1248x704 的
|
||
// → resizeCover 压缩覆盖写（原子写）。幂等：已是目标尺寸的行跳过，新库空跑。
|
||
func (s *datasetService) CompressExistingCovers(ctx context.Context) error {
|
||
list, err := dao.Dataset.ListAll(ctx)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
compressed := 0
|
||
for _, d := range list {
|
||
if d.Cover == "" || !isCoverName(d.Cover) {
|
||
continue
|
||
}
|
||
path := filepath.Join(common.DatasetImagesDir(ctx, d.Name), d.Cover)
|
||
data, err := os.ReadFile(path)
|
||
if err != nil {
|
||
continue
|
||
}
|
||
img, _, err := image.Decode(bytes.NewReader(data))
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "数据集 %s 封面解码失败，跳过压缩: %v", d.Name, err)
|
||
continue
|
||
}
|
||
if img.Bounds().Dx() == coverWidth && img.Bounds().Dy() == coverHeight {
|
||
continue
|
||
}
|
||
out, err := resizeCover(data)
|
||
if err != nil {
|
||
g.Log().Errorf(ctx, "数据集 %s 封面压缩失败: %+v", d.Name, err)
|
||
continue
|
||
}
|
||
if err := common.WriteFileAtomic(path, out); err != nil {
|
||
g.Log().Errorf(ctx, "数据集 %s 封面压缩写回失败: %+v", d.Name, err)
|
||
continue
|
||
}
|
||
compressed++
|
||
g.Log().Infof(ctx, "数据集 %s 封面压缩: %dx%d → %dx%d", d.Name, img.Bounds().Dx(), img.Bounds().Dy(), coverWidth, coverHeight)
|
||
}
|
||
if compressed > 0 {
|
||
g.Log().Infof(ctx, "封面存量压缩完成: 共 %d 个数据集", compressed)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// AdminCreateDataset 新建数据集：名称唯一（UNIQUE 兜底）+ 创建图片目录 +
|
||
// VLM 同步生成生成参数池（失败不阻断创建，poolsGenerated/poolError 反馈）。
|
||
func (s *datasetService) AdminCreateDataset(ctx context.Context, req *dto.AdminDatasetCreateReq) (*dto.AdminDatasetCreateRes, error) {
|
||
if req.Name == consts.NegativeDatasetName {
|
||
return nil, gerror.New("该名称为负样本库保留名，请更换数据集名称")
|
||
}
|
||
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")
|
||
}
|
||
sortOrder := int64(0)
|
||
if req.SortOrder != nil {
|
||
sortOrder = *req.SortOrder
|
||
}
|
||
id, err = dao.Dataset.Insert(ctx, &entity.Dataset{
|
||
Name: req.Name,
|
||
Source: req.Source,
|
||
NamePrefix: prefix,
|
||
SortOrder: sortOrder,
|
||
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, "创建数据集目录失败")
|
||
}
|
||
res := &dto.AdminDatasetCreateRes{Id: id}
|
||
// VLM 生成参数池：失败不阻断创建（返回原因，可用 gen-pools 补生成/手填提示词）
|
||
var species string
|
||
if pools, pErr := s.genPoolsWithVLM(ctx, req.Name); pErr != nil {
|
||
res.PoolError = pErr.Error()
|
||
} else if pools != nil {
|
||
if uErr := common.Serial().Submit(ctx, func() error {
|
||
return dao.Dataset.UpdatePools(ctx, id, pools)
|
||
}); uErr == nil {
|
||
res.PoolsGenerated = true
|
||
species = pools.GenSpecies
|
||
} else {
|
||
res.PoolError = "参数写入失败: " + uErr.Error()
|
||
}
|
||
}
|
||
// 封面：新建对话框预生成（req.Cover 非空且文件已落盘）优先——跳过自动生成；否则自动生成
|
||
// （16:9、1雄1雌），失败不阻断创建（前端提示可编辑模式重新生成）
|
||
if cover := strings.TrimSpace(req.Cover); cover != "" && isCoverName(cover) {
|
||
if _, statErr := os.Stat(filepath.Join(common.DatasetImagesDir(ctx, req.Name), filepath.Base(cover))); statErr == nil {
|
||
if uErr := common.Serial().Submit(ctx, func() error {
|
||
return dao.Dataset.UpdateConfigs(ctx, id, &entity.Dataset{Cover: filepath.Base(cover)})
|
||
}); uErr == nil {
|
||
res.CoverGenerated = true
|
||
return res, nil
|
||
} else {
|
||
res.CoverError = "封面写入失败: " + uErr.Error()
|
||
}
|
||
}
|
||
}
|
||
if _, coverErr := s.genCoverWithImageGen(ctx, id, req.Name, species); coverErr != nil {
|
||
res.CoverError = coverErr.Error()
|
||
} else {
|
||
res.CoverGenerated = true
|
||
}
|
||
return res, nil
|
||
}
|
||
|
||
// AdminNegativeLibrary 获取负样本库（技术设计.md「负样本库」）：source=negative 的固定保留名
|
||
// 特殊数据集，不存在则 Serial 内懒创建（防并发重建）+ 建目录；返回 id 与图片数，
|
||
// 管理端「负样本」tab 复用既有上传/图片列表/删除接口（负样本无标注/审核/清洗/生成流程）
|
||
func (s *datasetService) AdminNegativeLibrary(ctx context.Context, req *dto.AdminDatasetNegativeReq) (*dto.AdminDatasetNegativeRes, error) {
|
||
now := gtime.Now()
|
||
var lib *entity.Dataset
|
||
err := common.Serial().Submit(ctx, func() error {
|
||
existing, err := dao.Dataset.GetByName(ctx, consts.NegativeDatasetName)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if existing != nil {
|
||
lib = existing
|
||
return nil
|
||
}
|
||
id, err := dao.Dataset.Insert(ctx, &entity.Dataset{
|
||
Name: consts.NegativeDatasetName,
|
||
Source: consts.DatasetSourceNegative,
|
||
Status: consts.DatasetStatusBuilding,
|
||
CreatedAt: now,
|
||
UpdatedAt: now,
|
||
})
|
||
if err != nil {
|
||
return err
|
||
}
|
||
lib, err = dao.Dataset.GetById(ctx, id)
|
||
return err
|
||
})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if err := os.MkdirAll(common.DatasetImagesDir(ctx, lib.Name), 0o755); err != nil {
|
||
return nil, gerror.Wrap(err, "创建负样本库目录失败")
|
||
}
|
||
cnt, err := dao.DatasetImage.CountByDataset(ctx, lib.Id)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return &dto.AdminDatasetNegativeRes{Id: lib.Id, Name: lib.Name, ImageCount: cnt}, nil
|
||
}
|
||
|
||
// ensureJpeg 模型生成的图片统一转 jpg：解码校验 + 非 jpeg（png/webp 等）转 jpeg(Quality 92)。
|
||
// 生成服务返回格式不可控，扩展名统一 .jpg，内容须与扩展名一致（RF-DETR 按扩展名推断 mime 提交）。
|
||
func ensureJpeg(data []byte) ([]byte, error) {
|
||
img, format, err := image.Decode(bytes.NewReader(data))
|
||
if err != nil {
|
||
return nil, gerror.Wrap(err, "生成图片解码失败")
|
||
}
|
||
if format == "jpeg" {
|
||
return data, nil
|
||
}
|
||
var buf bytes.Buffer
|
||
if err := jpeg.Encode(&buf, img, &jpeg.Options{Quality: 92}); err != nil {
|
||
return nil, gerror.Wrap(err, "生成图片转 jpg 失败")
|
||
}
|
||
return buf.Bytes(), nil
|
||
}
|
||
|
||
// 封面统一尺寸（2026-08-31 用户定案）：生成/上传/存量压缩一律 1248x704（16:9 横幅，App 卡片展示规格）
|
||
const (
|
||
coverWidth = 1248
|
||
coverHeight = 704
|
||
)
|
||
|
||
// resizeCover 封面图片统一压缩为 1248x704：解码校验 → 按目标比例中心裁剪 → CatmullRom 缩放 → jpg(Quality 92)。
|
||
// 透明通道（png）先合成白底再缩放（jpg 无 alpha，白底比黑底自然）。
|
||
func resizeCover(data []byte) ([]byte, error) {
|
||
img, _, err := image.Decode(bytes.NewReader(data))
|
||
if err != nil {
|
||
return nil, gerror.Wrap(err, "封面图片解码失败")
|
||
}
|
||
src := image.NewRGBA(img.Bounds())
|
||
draw.Draw(src, src.Bounds(), image.White, image.Point{}, draw.Src)
|
||
draw.Draw(src, src.Bounds(), img, image.Point{}, draw.Over)
|
||
// 中心裁剪到 1248:704 比例（源图过宽裁左右，过窄裁上下）
|
||
sw, sh := src.Bounds().Dx(), src.Bounds().Dy()
|
||
targetRatio := float64(coverWidth) / float64(coverHeight)
|
||
if ratio := float64(sw) / float64(sh); ratio > targetRatio {
|
||
cw := int(float64(sh) * targetRatio)
|
||
src = cropCenter(src, cw, sh)
|
||
} else if ratio < targetRatio {
|
||
ch := int(float64(sw) / targetRatio)
|
||
src = cropCenter(src, sw, ch)
|
||
}
|
||
dst := image.NewRGBA(image.Rect(0, 0, coverWidth, coverHeight))
|
||
draw.CatmullRom.Scale(dst, dst.Bounds(), src, src.Bounds(), draw.Over, nil)
|
||
var buf bytes.Buffer
|
||
if err := jpeg.Encode(&buf, dst, &jpeg.Options{Quality: 92}); err != nil {
|
||
return nil, gerror.Wrap(err, "封面转 jpg 失败")
|
||
}
|
||
return buf.Bytes(), nil
|
||
}
|
||
|
||
// cropCenter 取源图中心 w×h 区域（src 为 RGBA，SubImage 返回同类型）
|
||
func cropCenter(src *image.RGBA, w, h int) *image.RGBA {
|
||
b := src.Bounds()
|
||
x0 := b.Min.X + (b.Dx()-w)/2
|
||
y0 := b.Min.Y + (b.Dy()-h)/2
|
||
return src.SubImage(image.Rect(x0, y0, x0+w, y0+h)).(*image.RGBA)
|
||
}
|
||
|
||
// genCoverWithImageGen 用 imageGen provider 生成数据集封面：16:9（1248x704）、画面 1 雄 1 雌，
|
||
// 同一物种雄性+雌性各一只；UUID 命名落盘，返回新文件名。
|
||
// datasetId>0：删旧封面文件 + 更新库 cover 字段（与手动上传同路径）；
|
||
// datasetId<=0：仅落盘（新建对话框预生成，数据集未创建，创建请求带 cover 回传写库）。
|
||
// 封面仅供展示，不入 dataset_image 表（不参与训练）；失败保留旧封面。
|
||
func (s *datasetService) genCoverWithImageGen(ctx context.Context, datasetId int64, datasetName, species string) (string, error) {
|
||
provider := common.ImageGen(ctx)
|
||
if provider == nil {
|
||
return "", gerror.New("图像生成服务未配置，未生成封面")
|
||
}
|
||
if species == "" {
|
||
species = datasetName
|
||
}
|
||
// 生成任务进行中直接调用：封面与任务图片同为 z-image 单张请求，由 local-ai 服务端排队串行，
|
||
// 无需整任务互斥（曾按 ListRunning 整任务拒绝，2026-09-02 放开）
|
||
prompt := fmt.Sprintf("野外实拍照片：一只雄性%s和一只雌性%s并排站立在开阔的自然栖息地中，"+
|
||
"雄雌各一只清晰可见，16:9 横幅构图，真实照片质感，光线自然，画面清晰美观", species, species)
|
||
// 180s：放开互斥后须覆盖 local-ai 排队等待（排队 ≤ 当前一张 ~40s + 自身生成），曾 120s 仅直发余量
|
||
genCtx, cancel := context.WithTimeout(ctx, 180*time.Second)
|
||
defer cancel()
|
||
// 直接指定封面统一尺寸 1248x704（z-image 接受 16 整除尺寸）；resizeCover 兜底统一规格
|
||
data, err := provider.Generate(genCtx, prompt, "1248x704")
|
||
if err != nil {
|
||
return "", gerror.Wrap(err, "封面生成失败")
|
||
}
|
||
if data, err = resizeCover(data); err != nil {
|
||
return "", err
|
||
}
|
||
dir := common.DatasetImagesDir(ctx, datasetName)
|
||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||
return "", gerror.Wrap(err, "创建图片目录失败")
|
||
}
|
||
if datasetId > 0 {
|
||
if dataset, dErr := dao.Dataset.GetById(ctx, datasetId); dErr == nil && dataset != nil && dataset.Cover != "" {
|
||
_ = os.Remove(filepath.Join(dir, filepath.Base(dataset.Cover)))
|
||
}
|
||
}
|
||
name := common.UuidV4() + ".jpg"
|
||
if err := common.WriteFileAtomic(filepath.Join(dir, name), data); err != nil {
|
||
return "", gerror.Wrap(err, "保存封面失败")
|
||
}
|
||
if datasetId > 0 {
|
||
if err := dao.Dataset.UpdateConfigs(ctx, datasetId, &entity.Dataset{Cover: name}); err != nil {
|
||
_ = os.Remove(filepath.Join(dir, name))
|
||
return "", err
|
||
}
|
||
}
|
||
return name, nil
|
||
}
|
||
|
||
// AdminGenPools 重新生成生成参数池（VLM 覆盖写；失败报错保留旧值——编辑对话框按钮入口）
|
||
func (s *datasetService) AdminGenPools(ctx context.Context, req *dto.AdminGenPoolsReq) (*dto.AdminGenPoolsRes, error) {
|
||
dataset, err := dao.Dataset.GetById(ctx, req.DatasetId)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if dataset == nil {
|
||
return nil, gerror.NewCode(common.CodeDatasetNotFound)
|
||
}
|
||
pools, err := s.genPoolsWithVLM(ctx, dataset.Name)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if pools == nil {
|
||
return nil, gerror.New("生成参数生成失败，请稍后重试")
|
||
}
|
||
if err := common.Serial().Submit(ctx, func() error {
|
||
return dao.Dataset.UpdatePools(ctx, dataset.Id, pools)
|
||
}); err != nil {
|
||
return nil, err
|
||
}
|
||
return &dto.AdminGenPoolsRes{
|
||
GenSpecies: pools.GenSpecies, GenTone: pools.GenTone, GenHeights: pools.GenHeights,
|
||
GenScenes: pools.GenScenes, GenActions: pools.GenActions, GenOcclusions: pools.GenOcclusions,
|
||
GenClasses: pools.GenClasses,
|
||
}, nil
|
||
}
|
||
|
||
// genPoolsWithVLM 调 qwen3.6-35b-a3b 生成数据集生成参数池（单物种规则：物种=数据集名，
|
||
// VLM 只生成轮廓色/站高/场景/动作/遮挡池）。
|
||
// 无 localAi 配置/JSON 非法/校验不过 → 返回 error（调用方决定是否阻断）；
|
||
// 生成任务进行中直接调用——VLM 与 z-image 同机共存，请求由 local-ai 服务端排队
|
||
// （曾按 ListRunning 整任务拒绝，2026-09-02 放开）；
|
||
// 校验通过返回 entity.Dataset（仅 7 个池字段，其余空）。
|
||
// 纯文本场景下 llama.cpp mmproj 需图片输入，传 64x64 纯灰占位图（提示词声明忽略图片）。
|
||
func (s *datasetService) genPoolsWithVLM(ctx context.Context, datasetName string) (*entity.Dataset, error) {
|
||
if common.LocalAiClient(ctx) == nil {
|
||
return nil, nil // 未配置标注服务：不生成也不报错（池为空，生成图片时提示补参数）
|
||
}
|
||
target := datasetName
|
||
prompt := fmt.Sprintf(
|
||
"你是野外野生动物监测数据集专家。请为数据集「%s」设计 AI 生成训练图的生成参数池。目标动物：%s。\n"+
|
||
"生成图规格：手机远距离拍摄的野外大远景实拍图，每张图有且只有一只目标动物，目标在画面中很小。\n"+
|
||
"请输出严格 JSON（只输出 JSON 对象本身，不要任何其他文字或代码块标记）：\n"+
|
||
"{\n \"tone\": \"深色\",\n \"height_cm\": 30,\n \"scenes\": [\"场景描述\", ...],\n"+
|
||
" \"actions\": [\"动作词\", ...],\n \"occlusions\": [\"遮挡描述\", ...]\n}\n"+
|
||
"要求：tone 为该物种轮廓色词（绝大多数「深色」，白化/浅色个体配「浅色」）；"+
|
||
"height_cm 为该物种站高厘米（鸟=脚到头顶、兽=蹲坐/站立肩高，数值 10~200）；"+
|
||
"scenes 4~6 条野外开阔场景，每条含「延伸向远方/地平线/开阔」等拉开距离的元素，符合该动物栖息习性"+
|
||
"（雉鸡→农田灌丛、野鸭→水面芦苇、兔子→草坡荒地）；按该动物冬季习性配场景：冬季在本地活跃（留鸟/不南迁，如雉鸡、兔子）须含至少 1 条冬季积雪场景（如雪后田野），冬季南迁物种（如野鸭、鹌鹑）不配雪地；若该物种会落树栖息（如雉鸡、锦鸡、马鸡、鸽子等树上夜栖鸟），场景须含林缘与树上落栖场景；不上树的地栖种（如石鸡、鹌鹑）不配树上场景；actions 4~6 条动作短语（如「觅食」「走动啄食」），不能以「在/于」开头；"+
|
||
"occlusions 4~6 条遮挡描述，覆盖轻（身体半掩）到重（只露头背），遮挡物符合该动物栖息环境（草丛/芦苇/树叶/枝条等）；具备树上/林缘场景的物种，遮挡池须含树叶/树枝类遮挡（不能只有草丛）。\n"+
|
||
"忽略图片内容，只按上述要求输出 JSON。",
|
||
datasetName, target)
|
||
// 64x64 纯灰占位图（VLM 输入占位，避免空白无法推理）
|
||
img := image.NewRGBA(image.Rect(0, 0, 64, 64))
|
||
var buf bytes.Buffer
|
||
if err := jpeg.Encode(&buf, img, &jpeg.Options{Quality: 60}); err != nil {
|
||
return nil, gerror.Wrap(err, "构造 VLM 输入失败")
|
||
}
|
||
// 120s：放开互斥后排队等待（生成任务图片 ~40s/张）+ 自身推理余量（曾 90s 仅直发余量）
|
||
genCtx, cancel := context.WithTimeout(ctx, 120*time.Second)
|
||
defer cancel()
|
||
content, err := common.QwenVL(genCtx, buf.Bytes(), "image/jpeg", prompt)
|
||
if err != nil {
|
||
return nil, gerror.Wrap(err, "VLM 生成参数失败")
|
||
}
|
||
return s.parseGenPoolsJSON(content, target)
|
||
}
|
||
|
||
// parseGenPoolsJSON 解析并校验 VLM 输出的参数池 JSON（容忍代码块/前后缀），
|
||
// 物种固定为 species（单物种规则：每数据集一个物种，第一标注类别名=物种，第二类别名="suspect"）。
|
||
// 校验失败返回 error（宁可报错也不落不合格池——生成图按池组装，池质量决定训练数据质量）。
|
||
func (s *datasetService) parseGenPoolsJSON(content string, species string) (*entity.Dataset, error) {
|
||
m := regexp.MustCompile(`\{[\s\S]*\}`).FindString(content)
|
||
if m == "" {
|
||
return nil, gerror.New("VLM 未输出有效 JSON")
|
||
}
|
||
var raw struct {
|
||
Tone string `json:"tone"`
|
||
HeightCm float64 `json:"height_cm"`
|
||
Scenes []string `json:"scenes"`
|
||
Actions []string `json:"actions"`
|
||
Occlusions []string `json:"occlusions"`
|
||
}
|
||
if err := json.Unmarshal([]byte(m), &raw); err != nil {
|
||
return nil, gerror.Wrap(err, "VLM 输出 JSON 解析失败")
|
||
}
|
||
valid := func(list []string, min int, maxLen int) bool {
|
||
if len(list) < min {
|
||
return false
|
||
}
|
||
for _, v := range list {
|
||
if strings.TrimSpace(v) == "" || len([]rune(v)) > maxLen {
|
||
return false
|
||
}
|
||
}
|
||
return true
|
||
}
|
||
if !valid(raw.Scenes, 3, 100) || !valid(raw.Actions, 3, 30) || !valid(raw.Occlusions, 3, 100) {
|
||
return nil, gerror.New("VLM 场景/动作/遮挡池非法（各需至少 3 条）")
|
||
}
|
||
tone := "深色"
|
||
if strings.TrimSpace(raw.Tone) != "" {
|
||
tone = strings.TrimSpace(raw.Tone)
|
||
}
|
||
// 站高存数据集表 gen_heights（数据库维护，用户可在编辑界面手改真实值），VLM 输出为初始值
|
||
height := 35.0
|
||
if raw.HeightCm >= 10 && raw.HeightCm <= 200 {
|
||
height = raw.HeightCm
|
||
}
|
||
scenesRaw, _ := json.Marshal(raw.Scenes)
|
||
actionsRaw, _ := json.Marshal(raw.Actions)
|
||
occlRaw, _ := json.Marshal(raw.Occlusions)
|
||
// 单值列：物种 / 轮廓色词 / 站高 / 第二类别名（第一类别名=物种，训练 data.yaml names / 模型 labels）
|
||
return &entity.Dataset{
|
||
GenSpecies: species,
|
||
GenTone: tone,
|
||
GenHeights: height,
|
||
GenScenes: string(scenesRaw),
|
||
GenActions: string(actionsRaw),
|
||
GenOcclusions: string(occlRaw),
|
||
GenClasses: "suspect",
|
||
}, nil
|
||
}
|
||
|
||
// AdminDeleteDataset 删除数据集：有 running 标注任务 / 该数据集训练进行中 / 已发布模型版本时拒绝
|
||
// （训练产物与模型为付费资产，需先删除模型版本再删数据集）；排队训练任务直接置 failed。
|
||
// 删除 = 删图片/模型目录 + 删记录（标注随图片行删除，Serial 单写者串行）。
|
||
func (s *datasetService) AdminDeleteDataset(ctx context.Context, req *dto.AdminDatasetDeleteReq) (*dto.AdminDatasetDeleteRes, error) {
|
||
var name, modelName 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)
|
||
}
|
||
if d.Source == consts.DatasetSourceNegative {
|
||
return gerror.New("负样本库不可删除整库（如需清理请逐张删除图片）")
|
||
}
|
||
name = d.Name
|
||
modelName = modelFileName(d.Name, d.NamePrefix)
|
||
// 预标注任务进行中（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("该数据集有训练任务进行中，无法删除")
|
||
}
|
||
// 排队任务引用将删的目录，晋级必失败：直接置 failed 带出原因（双档位串行队列 2026-09-03 起支持排队）
|
||
if err := dao.Training.FailQueuedByDataset(ctx, d.Id, "数据集已删除，排队训练取消"); err != nil {
|
||
return err
|
||
}
|
||
// 模型版本记录随数据集级联删除（管理端无模型管理界面，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
|
||
}
|
||
// 文件清理（图片目录 + 当前生效模型文件，s/n 双档位文件一并清理；删除失败仅记日志，记录已删）
|
||
paths := []string{
|
||
common.DatasetImagesDir(ctx, name),
|
||
common.TrainingModelPath(ctx, modelName),
|
||
common.TrainingModelPath(ctx, modelName+consts.TrainingVariantNFileSuffix),
|
||
}
|
||
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
|
||
}
|
||
// 标注不再强制（2026-09-04 自动标注退场）：未配置 localAi 也可上传，图片以未标注态入池
|
||
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
|
||
}
|
||
// 上传不触发标注（2026-09-04 自动标注退场）：图片以未标注态入池（管理端预标/众包任务池）
|
||
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
|
||
}
|
||
|
||
// singleAnimalClause 目标数量约束：每张仅一只动物，手填提示词统一追加（模板已由 {count} 固定 "1只"）
|
||
const singleAnimalClause = "画面中有且只有这一只动物，没有任何其他动物"
|
||
|
||
// buildPromptFromTemplate 按 config 通用模板组装提示词（单物种规则：物种固定取数据集 gen_species，
|
||
// 空则数据集名本身——生成表单无物种输入，不随机）。
|
||
// 物种只写名字不写羽毛细节（模型对正确名称自带外观先验）；场景/动作/遮挡从数据集表池随机
|
||
// （创建时 VLM 生成，config.yml 不再兜底；池空报错——生成图按池组装，池质量决定训练数据质量）；
|
||
// 场景池条目兼容捆绑组（树栖物种场景与动作/遮挡联动，见 parseScenePool）；光线走 config 通用池 lights。
|
||
func (s *datasetService) buildPromptFromTemplate(ctx context.Context, dataset *entity.Dataset, dist int, size string) (string, error) {
|
||
tpl := g.Cfg().MustGet(ctx, "imageGen.promptTemplate").String()
|
||
if tpl == "" {
|
||
return "", gerror.New("未配置通用提示词模板(imageGen.promptTemplate)，请手填提示词")
|
||
}
|
||
species := strings.TrimSpace(dataset.GenSpecies)
|
||
if species == "" {
|
||
species = dataset.Name
|
||
}
|
||
// 场景/动作/遮挡按数据集习性取池（表存储，VLM 生成；池空报错——生成图按池组装，池质量决定训练数据质量）。
|
||
// 场景池条目兼容捆绑组对象：选中捆绑组时动作/遮挡仅从组内池随机（生境联动），否则走全局池
|
||
poolOf := func(raw string) []string {
|
||
var pool []string
|
||
if json.Unmarshal([]byte(raw), &pool) != nil {
|
||
return nil
|
||
}
|
||
return pool
|
||
}
|
||
scene, action, occlusion := "", "", ""
|
||
bundle := pickSceneEntry(parseScenePool(dataset.GenScenes))
|
||
if bundle != nil {
|
||
scene = bundle.Scene
|
||
if len(bundle.Actions) > 0 {
|
||
action = pickCfgList(bundle.Actions)
|
||
}
|
||
if len(bundle.Occlusions) > 0 {
|
||
occlusion = pickCfgList(bundle.Occlusions)
|
||
}
|
||
}
|
||
if action == "" {
|
||
action = pickCfgList(poolOf(dataset.GenActions))
|
||
}
|
||
if occlusion == "" {
|
||
occlusion = pickCfgList(poolOf(dataset.GenOcclusions))
|
||
}
|
||
if scene == "" || action == "" || occlusion == "" {
|
||
return "", gerror.New("数据集「" + dataset.Name + "」未配置生成参数池（场景/动作/遮挡），请在编辑数据集生成或手填提示词")
|
||
}
|
||
light := pickCfg(ctx, "imageGen.lights")
|
||
// 轮廓色词（表 gen_tone，白化个体如白雉鸡配浅色，未配置默认深色剪影）
|
||
tone := strings.TrimSpace(dataset.GenTone)
|
||
if tone == "" {
|
||
tone = "深色"
|
||
}
|
||
// 距离描述从表单固定值生成（25 → "25米外"）
|
||
distWord := fmt.Sprintf("%d米外", dist)
|
||
// 位置锚 {position}：默认地平线锚（地面/水面场景，目标在地平线附近语义成立）；
|
||
// 捆绑组 position 覆盖（树栖场景目标须锚到枝头/树冠，否则「地平线附近有且仅有斑鸠」仍落在地面）
|
||
position := "在" + distWord + "的地平线附近"
|
||
if bundle != nil && strings.TrimSpace(bundle.Position) != "" {
|
||
position = strings.ReplaceAll(bundle.Position, "{distanceWord}", distWord)
|
||
}
|
||
// 尺寸提示：按物理公式把表单距离换算成目标在图中的像素高与画面占比
|
||
sizeHint, pctVal, sErr := s.buildSizeHint(ctx, dataset, dist, size)
|
||
if sErr != nil {
|
||
return "", sErr
|
||
}
|
||
// 占比 <2% 时切换远景小目标模板：目标小到动作/遮挡/性别细节无法呈现，
|
||
// 保留细节描写会迫使模型把目标画大（2026-08-31 实测两次失败后定案）
|
||
if pctVal < 2.0 {
|
||
if tiny := g.Cfg().MustGet(ctx, "imageGen.promptTemplateTiny").String(); tiny != "" {
|
||
tpl = tiny
|
||
}
|
||
}
|
||
// 每张固定 1 个目标（数量词固定，animal_count 列同为 1，仅信息性存储；
|
||
// 生成图可能实际含多个目标，预标检出超 3 个按置信度取前 3，2026-09-02）
|
||
countWord := "1只"
|
||
// 性别随机（50/50）：两性体型与外观差异大（雉鸡雄艳雌褐、野鸭雄艳雌素等），
|
||
// 每张随机让训练数据覆盖两性形态
|
||
sexWord := "雄性"
|
||
if rand.Intn(2) == 1 {
|
||
sexWord = "雌性"
|
||
}
|
||
out := tpl
|
||
for k, v := range map[string]string{
|
||
"{scene}": scene, "{species}": sexWord + species, "{count}": countWord,
|
||
"{action}": action, "{occlusion}": occlusion, "{light}": light,
|
||
"{distanceWord}": distWord, "{sizeHint}": sizeHint, "{tone}": tone,
|
||
"{position}": position,
|
||
} {
|
||
out = strings.ReplaceAll(out, k, v)
|
||
}
|
||
g.Log().Infof(ctx, "模板组装提示词: %s", out)
|
||
return out, nil
|
||
}
|
||
|
||
// buildSizeHint 尺寸提示：目标像素高 = 体高 × focalPx ÷ 距离（focalPx = 图高/2 ÷ tan(vfov/2)），
|
||
// 尺寸词按「像素高÷图高」画面占比分级（量化像素数对文生图模型无意义；固定物体类比词
|
||
// 不跨物种/距离通用，2026-08-31 用户否决），输出如「（在画面中很小的身影，约占画面高度1.0%）」；
|
||
// 返回占比供调用方按 <2% 切换远景小目标模板。占比保留一位小数，避免出现 0%
|
||
func (s *datasetService) buildSizeHint(ctx context.Context, dataset *entity.Dataset, dist int, size string) (string, float64, error) {
|
||
_, imgH, err := parseGenSize(size)
|
||
if err != nil {
|
||
return "", 0, err
|
||
}
|
||
// 站高 cm 从数据集表 gen_heights 读（数据库维护，编辑界面手改真实值；<=0 物理兜底 35cm），换算成米参与投影公式
|
||
heightCM := dataset.GenHeights
|
||
if heightCM <= 0 {
|
||
heightCM = 35.0
|
||
}
|
||
heightM := heightCM / 100
|
||
// 镜头固定为手机默认主摄（垂直视场角配置值，默认 52°；2026-08-28 曾随机 [15,90]、曾试超广角，用户最终定为默认镜头）
|
||
vfov := g.Cfg().MustGet(ctx, "imageGen.assumedVfovDeg", 52.0).Float64() * math.Pi / 180
|
||
focalPx := float64(imgH) / 2 / math.Tan(vfov/2)
|
||
pxAt := func(dist int) int {
|
||
if dist <= 0 {
|
||
dist = 1
|
||
}
|
||
return int(math.Round(heightM * focalPx / float64(dist)))
|
||
}
|
||
pct := func(px int) string {
|
||
return fmt.Sprintf("%.1f", float64(px)*100/float64(imgH))
|
||
}
|
||
// 尺寸词按画面占比分级（占比=像素高÷图高，与物种/距离无关：同占比同观感；
|
||
// 固定物体类比词（针尖/米粒/拳头）不跨物种通用，用户否决 2026-08-31；
|
||
// 站高 0.6m、focalPx≈1280 时：60m≈13px≈1.0%、25m≈31px≈2.5%、100m≈8px≈0.6%）
|
||
px := pxAt(dist)
|
||
pctVal := float64(px) * 100 / float64(imgH)
|
||
word := "中等大小的身影"
|
||
switch {
|
||
case pctVal < 1.0:
|
||
word = "微小到几乎难以辨认的身影"
|
||
case pctVal < 2.0:
|
||
word = "很小的身影"
|
||
case pctVal < 5.0:
|
||
word = "较小的身影"
|
||
case pctVal < 10.0:
|
||
word = "中等大小的身影"
|
||
default:
|
||
word = "较大的身影"
|
||
}
|
||
return fmt.Sprintf("（在画面中%s，约占画面高度%s%%）", word, pct(px)), pctVal, 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
|
||
}
|
||
|
||
// pickCfg 取字符串数组配置项随机一项（空池返回空串）
|
||
func pickCfg(ctx context.Context, key string) string {
|
||
return pickCfgList(g.Cfg().MustGet(ctx, key).Strings())
|
||
}
|
||
|
||
// pickCfgList 从字符串数组随机取一项（空池返回空串）
|
||
func pickCfgList(pool []string) string {
|
||
if len(pool) == 0 {
|
||
return ""
|
||
}
|
||
return pool[rand.Intn(len(pool))]
|
||
}
|
||
|
||
// genSceneEntry 场景池条目：字符串条目解析为仅含 Scene 的本结构（weight=1，动作/遮挡走全局池）；
|
||
// 捆绑组对象 {scene, weight, actions, occlusions, position} 的动作/遮挡与场景联动（2026-09-04，
|
||
// 斑鸠树栖 80%——三池独立随机会拼出「树上场景+地面动作/草丛遮挡」矛盾提示词，见技术设计.md）；
|
||
// position 覆盖默认位置锚「在N米外的地平线附近」（树栖目标锚地平线=语义落地，须锚到枝头，
|
||
// 内可用 {distanceWord} 占位）
|
||
type genSceneEntry struct {
|
||
Scene string `json:"scene"`
|
||
Weight int `json:"weight"`
|
||
Actions []string `json:"actions"`
|
||
Occlusions []string `json:"occlusions"`
|
||
Position string `json:"position"`
|
||
}
|
||
|
||
// parseScenePool 解析场景池 JSON 数组：条目兼容字符串与捆绑组对象，非法条目跳过
|
||
func parseScenePool(raw string) []genSceneEntry {
|
||
var items []json.RawMessage
|
||
if json.Unmarshal([]byte(raw), &items) != nil {
|
||
return nil
|
||
}
|
||
entries := make([]genSceneEntry, 0, len(items))
|
||
for _, it := range items {
|
||
var s string
|
||
if json.Unmarshal(it, &s) == nil {
|
||
if strings.TrimSpace(s) == "" {
|
||
continue // 空串/null 条目（Unmarshal null 入 string 成功但为空）不进池
|
||
}
|
||
entries = append(entries, genSceneEntry{Scene: s, Weight: 1})
|
||
continue
|
||
}
|
||
var b genSceneEntry
|
||
if json.Unmarshal(it, &b) != nil || strings.TrimSpace(b.Scene) == "" {
|
||
continue
|
||
}
|
||
if b.Weight < 1 {
|
||
b.Weight = 1
|
||
}
|
||
entries = append(entries, b)
|
||
}
|
||
return entries
|
||
}
|
||
|
||
// pickSceneEntry 场景池加权随机一项（weight 缺省 1；空池返回 nil）
|
||
func pickSceneEntry(pool []genSceneEntry) *genSceneEntry {
|
||
total := 0
|
||
for i := range pool {
|
||
total += pool[i].Weight
|
||
}
|
||
if total <= 0 {
|
||
return nil
|
||
}
|
||
n := rand.Intn(total)
|
||
for i := range pool {
|
||
n -= pool[i].Weight
|
||
if n < 0 {
|
||
return &pool[i]
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// AdminGenerateImages AI 生成图片（异步任务）：校验通过后插 gen_task 立即返回 TaskId，
|
||
// 后台协程逐张生成（不设调用超时，失败由 provider 返回错误决定），进度落库供前端轮询；
|
||
// 完成（或部分失败）后对本次新增图触发自动标注（付费资产，失败保留已生成图不删除）。
|
||
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 dataset.Source == consts.DatasetSourceNegative {
|
||
return nil, gerror.New("负样本库不支持 AI 生成图片")
|
||
}
|
||
// 标注不再强制（2026-09-04 自动标注退场）：未配置 localAi 也可生成，图以未标注态入池
|
||
dir := common.DatasetImagesDir(ctx, dataset.Name)
|
||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||
return nil, gerror.Wrap(err, "创建图片目录失败")
|
||
}
|
||
// 命名：数据集属性 name_prefix（创建/编辑表单维护）非空时按 <前缀>_<两位序号>.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
|
||
}
|
||
}
|
||
// Serial 内并发检查 + 插任务：同数据集已有 running 生成任务则拒绝
|
||
now := gtime.Now()
|
||
var taskId int64
|
||
if err := common.Serial().Submit(ctx, func() error {
|
||
running, err := dao.GenTask.GetRunningByDataset(ctx, dataset.Id)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if running != nil {
|
||
return gerror.NewCode(common.CodeGenTaskRunning)
|
||
}
|
||
taskId, err = dao.GenTask.Insert(ctx, &entity.GenTask{
|
||
DatasetId: dataset.Id,
|
||
Status: consts.GenTaskRunning,
|
||
Total: req.Count,
|
||
CreatedAt: now,
|
||
})
|
||
return err
|
||
}); err != nil {
|
||
return nil, err
|
||
}
|
||
s.runGenTask(taskId, dataset, req, provider, now, prefix, nextIdx)
|
||
return &dto.AdminDatasetGenerateRes{TaskId: taskId, Total: req.Count}, nil
|
||
}
|
||
|
||
// recoverGenTasks 服务重启恢复：孤儿 running 生成任务置 failed + 错误提示
|
||
// （生成中断后已生成图为付费资产保留；重新发起即可补齐，startDetection 并发检查不受孤儿影响）。
|
||
func (s *datasetService) recoverGenTasks(ctx context.Context) {
|
||
list, err := dao.GenTask.ListRunning(ctx)
|
||
if err != nil {
|
||
g.Log().Errorf(ctx, "恢复生成任务失败: %+v", err)
|
||
return
|
||
}
|
||
for _, t := range list {
|
||
if err := dao.GenTask.Finish(ctx, t.Id, "服务重启，任务中断，可重新发起"); err != nil {
|
||
g.Log().Errorf(ctx, "恢复生成任务 %d 失败: %+v", t.Id, err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// runGenTask 生成执行协程（生命周期任务）：脱离请求 ctx（请求结束即取消会让 Submit 秒退 +
|
||
// Finish 静默失败 → 任务悬挂），池内逐张生成（imageGen.poolSize 并发，z-image 显存独占默认 1），
|
||
// 每张成功后 Serial 更新进度；全部完成 → 更新数据集计数 → 自动标注 → 置 done；
|
||
// 中途失败 → 已生成图先触发自动标注 → 置 failed（错误信息带已生成张数）。
|
||
func (s *datasetService) runGenTask(taskId int64, dataset *entity.Dataset, req *dto.AdminDatasetGenerateReq, provider common.ImageGenProvider, now *gtime.Time, prefix string, nextIdx int) {
|
||
bgCtx := context.Background()
|
||
go func() {
|
||
dir := common.DatasetImagesDir(bgCtx, dataset.Name)
|
||
// 异步任务不设调用超时：超时只放弃等待、取消不了 local-ai backend，
|
||
// 会造成「任务已失败、模型仍在跑」的假失败并占住 poolSize=1 的池；失败判定交给 provider 真实返回
|
||
handWritten := strings.TrimSpace(req.Prompt) != ""
|
||
generated := 0
|
||
addedIds := make([]int64, 0, req.Count)
|
||
failed := ""
|
||
for i := 0; i < req.Count; i++ {
|
||
// 提示词：手填优先（统一追加数量约束：仅一只动物，拍摄角度不带，
|
||
// 交由模型平视先验自由构图）；未手填走 config 通用模板逐张组装。距离校验已取消（2026-08-28）
|
||
prompt := strings.TrimSpace(req.Prompt)
|
||
if handWritten {
|
||
prompt = strings.TrimSuffix(prompt, "。") + "。"
|
||
if !strings.Contains(prompt, "一只") && !strings.Contains(prompt, "1只") && !strings.Contains(prompt, "唯一") {
|
||
prompt += singleAnimalClause
|
||
}
|
||
} else {
|
||
var tErr error
|
||
if prompt, tErr = s.buildPromptFromTemplate(bgCtx, dataset, req.Distance, req.Size); tErr != nil {
|
||
failed = tErr.Error()
|
||
break
|
||
}
|
||
}
|
||
var data []byte
|
||
err := common.GenTaskPoolInstance().Submit(bgCtx, func(ctx context.Context) error {
|
||
d, e := provider.Generate(ctx, prompt, req.Size)
|
||
data = d
|
||
return e
|
||
})
|
||
if err != nil {
|
||
// 付费资产保留原则：已生成的不删除
|
||
failed = fmt.Sprintf("第 %d 张生成失败(已生成 %d 张): %v", i+1, generated, err)
|
||
break
|
||
}
|
||
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)
|
||
}
|
||
// 生成图统一转 jpg（模型可能返回 png/webp，扩展名 .jpg 须与内容一致）
|
||
data, jErr := ensureJpeg(data)
|
||
if jErr != nil {
|
||
failed = fmt.Sprintf("第 %d 张转 jpg 失败(已生成 %d 张): %v", i+1, generated, jErr)
|
||
break
|
||
}
|
||
if err := common.WriteFileAtomic(filepath.Join(dir, filename), data); err != nil {
|
||
failed = fmt.Sprintf("第 %d 张保存失败(已生成 %d 张): %v", i+1, generated, err)
|
||
break
|
||
}
|
||
var imgId int64
|
||
if insErr := common.Serial().Submit(bgCtx, func() error {
|
||
id, err := dao.DatasetImage.Insert(bgCtx, &entity.DatasetImage{
|
||
DatasetId: dataset.Id,
|
||
Filename: filename,
|
||
Source: "ai",
|
||
AnimalCount: 1,
|
||
CreatedAt: now,
|
||
})
|
||
imgId = id
|
||
return err
|
||
}); insErr != nil {
|
||
failed = fmt.Sprintf("第 %d 张入库失败(已生成 %d 张): %v", i+1, generated, insErr)
|
||
break
|
||
}
|
||
addedIds = append(addedIds, imgId)
|
||
generated++
|
||
_ = common.Serial().Submit(bgCtx, func() error {
|
||
return dao.GenTask.UpdateProgress(bgCtx, taskId, generated)
|
||
})
|
||
}
|
||
if failed != "" {
|
||
// 部分失败：已入库图按付费资产保留，计数须计入（曾漏更新致 image_count 漂移为负）
|
||
if err := common.Serial().Submit(bgCtx, func() error {
|
||
return dao.Dataset.UpdateCounters(bgCtx, dataset.Id, int64(generated), -1, "")
|
||
}); err != nil {
|
||
g.Log().Errorf(bgCtx, "生成任务 %d 部分失败后更新数据集计数失败: %+v", taskId, err)
|
||
}
|
||
_ = dao.GenTask.Finish(bgCtx, taskId, failed)
|
||
return
|
||
}
|
||
// 全部成功：更新数据集计数 → 置 done（图以未标注态入池，2026-09-04 自动标注退场）
|
||
if err := common.Serial().Submit(bgCtx, func() error {
|
||
return dao.Dataset.UpdateCounters(bgCtx, dataset.Id, int64(generated), -1, "")
|
||
}); err != nil {
|
||
g.Log().Errorf(bgCtx, "生成任务 %d 更新数据集计数失败: %+v", taskId, err)
|
||
}
|
||
_ = dao.GenTask.Finish(bgCtx, taskId, "")
|
||
}()
|
||
}
|
||
|
||
// AdminGenTaskQuery 生成任务进度（最近一次；无任务返回 nil）
|
||
func (s *datasetService) AdminGenTaskQuery(ctx context.Context, req *dto.AdminGenTaskQueryReq) (*dto.AdminGenTaskQueryRes, error) {
|
||
dataset, err := dao.Dataset.GetById(ctx, req.DatasetId)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if dataset == nil {
|
||
return nil, gerror.NewCode(common.CodeDatasetNotFound)
|
||
}
|
||
t, err := dao.GenTask.LatestByDataset(ctx, dataset.Id)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if t == nil {
|
||
return nil, nil
|
||
}
|
||
return &dto.AdminGenTaskQueryRes{
|
||
Id: t.Id,
|
||
DatasetId: t.DatasetId,
|
||
Status: t.Status,
|
||
Total: t.Total,
|
||
Done: t.Done,
|
||
Error: t.Error,
|
||
CreatedAt: t.CreatedAt,
|
||
FinishedAt: t.FinishedAt,
|
||
}, 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,
|
||
Url: datasetImageUrl(ctx, dataset.Id, v.Filename),
|
||
CleanExcluded: v.CleanExcluded,
|
||
AnnotateTaskId: v.AnnotateTaskId,
|
||
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
|
||
}
|
||
|
||
// ---------- 数据清洗（训练集去冗余；语义与档位见技术设计.md「数据清洗」/consts.CleanBuckets） ----------
|
||
|
||
// cleanImageStat 参与统计的单图（labels_json 含 class 0 框：面积最大者定档，其框高为展示主尺寸；
|
||
// 档内保留取舍看 minHPct——小目标样本稀缺，先保最小确认目标更小的图）
|
||
type cleanImageStat struct {
|
||
id int64
|
||
filename string
|
||
source string
|
||
area float64 // 最大 class0 框面积 w*h（找最大框用）
|
||
hPct float64 // 该框高占图高百分比（定档与主尺寸展示）
|
||
minHPct float64 // 图内最小 class0 框占比（优先保留依据；单目标图 = hPct）
|
||
targetCnt int // class0 目标总数
|
||
bucket int
|
||
}
|
||
|
||
// AdminCleanPreview 数据清洗预览：桶分布 + 超配桶候选清单 + 已排除清单。
|
||
// 候选生成不随机——超配桶内先按 minHPct 升序扫描（小目标样本稀缺先占保留名额，同尺寸按 id 稳定），
|
||
// 再整图 dHash 贪心：与已保留图最小汉明距离 > 阈值才保留（连拍帧哈希近距只留首帧），
|
||
// 留够配额即停，其余进候选；文件缺失/解码失败无法比较，宁可保留。
|
||
func (s *datasetService) AdminCleanPreview(ctx context.Context, req *dto.AdminDatasetCleanPreviewReq) (*dto.AdminDatasetCleanPreviewRes, error) {
|
||
dataset, err := dao.Dataset.GetById(ctx, req.DatasetId)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if dataset == nil {
|
||
return nil, gerror.NewCode(common.CodeDatasetNotFound)
|
||
}
|
||
quotas := make([]int, len(consts.CleanBuckets))
|
||
for i, b := range consts.CleanBuckets {
|
||
quotas[i] = b.Quota
|
||
}
|
||
if len(req.Quotas) > 0 {
|
||
if len(req.Quotas) != len(consts.CleanBuckets) {
|
||
return nil, gerror.Newf("配额数量须与档数一致（共 %d 档）", len(consts.CleanBuckets))
|
||
}
|
||
for i, q := range req.Quotas {
|
||
if q < 0 {
|
||
return nil, gerror.New("配额不能为负数")
|
||
}
|
||
if i == 0 {
|
||
continue // <2% 极远档豁免，配额固定不适用
|
||
}
|
||
quotas[i] = q
|
||
}
|
||
}
|
||
images, err := dao.DatasetImage.ListByDataset(ctx, dataset.Id)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
// 贪心扫描按 id 升序（ListByDataset 为文件名数值倒序，须重排保证确定性）
|
||
sort.Slice(images, func(i, j int) bool { return images[i].Id < images[j].Id })
|
||
res := &dto.AdminDatasetCleanPreviewRes{
|
||
Buckets: make([]*dto.AdminCleanBucketRes, len(consts.CleanBuckets)),
|
||
Candidates: []*dto.AdminCleanImage{},
|
||
ExcludedImages: []*dto.AdminCleanImage{},
|
||
}
|
||
for i, b := range consts.CleanBuckets {
|
||
res.Buckets[i] = &dto.AdminCleanBucketRes{Label: b.Label, Exempt: i == 0, Quota: quotas[i]}
|
||
}
|
||
// 第一遍：定档统计（仅含 class 0 的图），已排除清单全量收集（当前无 class 0 标注的也列出，恢复不丢）
|
||
actives := make([][]*cleanImageStat, len(consts.CleanBuckets)) // 桶内未排除图，保持 id 升序
|
||
for _, img := range images {
|
||
if img.CleanExcluded == 1 {
|
||
res.Excluded++
|
||
}
|
||
info := &dto.AdminCleanImage{ImageId: img.Id, Filename: img.Filename, Source: img.Source}
|
||
var stat = cleanImageStat{id: img.Id, filename: img.Filename, source: img.Source}
|
||
if img.LabelsJson != "" && img.LabelsJson != "[]" {
|
||
var boxes []*dto.AdminLabelBox
|
||
if json.Unmarshal([]byte(img.LabelsJson), &boxes) == nil {
|
||
for _, b := range boxes {
|
||
if b.Class != 0 || b.W <= 0 || b.H <= 0 {
|
||
continue // 只认确认框(class 0)；疑似框(class 1)不入尺寸统计
|
||
}
|
||
pct := b.H * 100
|
||
stat.targetCnt++
|
||
if stat.minHPct == 0 || pct < stat.minHPct {
|
||
stat.minHPct = pct
|
||
}
|
||
if b.W*b.H > stat.area {
|
||
stat.area = b.W * b.H
|
||
stat.hPct = pct
|
||
stat.bucket = cleanBucketIdx(pct)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
if stat.area <= 0 { // 无确认目标（未标注/空检出/仅 class 1）：不入档
|
||
if img.CleanExcluded == 1 {
|
||
res.ExcludedImages = append(res.ExcludedImages, info)
|
||
}
|
||
continue
|
||
}
|
||
info.BucketLabel = res.Buckets[stat.bucket].Label
|
||
info.BoxHeightPct = math.Round(stat.hPct*10) / 10
|
||
bk := res.Buckets[stat.bucket]
|
||
bk.Total++
|
||
res.Total++
|
||
if img.CleanExcluded == 1 {
|
||
bk.Excluded++
|
||
res.ExcludedImages = append(res.ExcludedImages, info)
|
||
} else {
|
||
actives[stat.bucket] = append(actives[stat.bucket], &stat)
|
||
}
|
||
}
|
||
// 第二遍：超配桶（未排除 > 配额）内 dHash 贪心挑保留，其余进候选
|
||
imgDir := common.DatasetImagesDir(ctx, dataset.Name)
|
||
for i, bk := range res.Buckets {
|
||
if bk.Exempt {
|
||
continue // 极远档豁免：固定保留，不进候选（over 恒 0）
|
||
}
|
||
over := len(actives[i]) - bk.Quota
|
||
if over <= 0 {
|
||
continue
|
||
}
|
||
bk.Over = over
|
||
// 小目标样本稀缺：同档内先保最小确认目标更小的图（同尺寸按 id 保持稳定输出）
|
||
sort.Slice(actives[i], func(a, b int) bool {
|
||
if actives[i][a].minHPct != actives[i][b].minHPct {
|
||
return actives[i][a].minHPct < actives[i][b].minHPct
|
||
}
|
||
return actives[i][a].id < actives[i][b].id
|
||
})
|
||
kept, keptHashes := 0, make([]uint64, 0, bk.Quota)
|
||
for _, a := range actives[i] {
|
||
cand := &dto.AdminCleanImage{ImageId: a.id, Filename: a.filename, Source: a.source,
|
||
BucketLabel: bk.Label, BoxHeightPct: math.Round(a.hPct*10) / 10,
|
||
MinHeightPct: math.Round(a.minHPct*10) / 10, TargetCount: a.targetCnt}
|
||
if kept >= bk.Quota { // 留够配额，其余全进候选
|
||
res.Candidates = append(res.Candidates, cand)
|
||
continue
|
||
}
|
||
h, ok := imageDHashFile(ctx, imgDir, a.filename)
|
||
if !ok { // 文件缺失/解码失败无法比较，宁可保留
|
||
kept++
|
||
continue
|
||
}
|
||
if minHamming(h, keptHashes) > consts.CleanHashHamming { // 与已保留图均拉开距离才留
|
||
keptHashes = append(keptHashes, h)
|
||
kept++
|
||
continue
|
||
}
|
||
res.Candidates = append(res.Candidates, cand)
|
||
}
|
||
}
|
||
return res, nil
|
||
}
|
||
|
||
// AdminCleanApply 数据清洗执行：批量置/清 clean_excluded（排除不删文件，只影响训练集打包；
|
||
// 恢复 = Exclude=false 重提该批 id）。imageIds 须全属该数据集。
|
||
func (s *datasetService) AdminCleanApply(ctx context.Context, req *dto.AdminDatasetCleanApplyReq) (*dto.AdminDatasetCleanApplyRes, error) {
|
||
dataset, err := dao.Dataset.GetById(ctx, req.DatasetId)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if dataset == nil {
|
||
return nil, gerror.NewCode(common.CodeDatasetNotFound)
|
||
}
|
||
ids := make([]int64, 0, len(req.ImageIds))
|
||
seen := make(map[int64]bool, len(req.ImageIds))
|
||
for _, id := range req.ImageIds {
|
||
if !seen[id] {
|
||
seen[id] = true
|
||
ids = append(ids, id)
|
||
}
|
||
}
|
||
images, err := dao.DatasetImage.GetByIds(ctx, ids)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if len(images) != len(ids) {
|
||
return nil, gerror.New("部分图片不存在，请刷新后重试")
|
||
}
|
||
for _, img := range images {
|
||
if img.DatasetId != dataset.Id {
|
||
return nil, gerror.Newf("图片 %s 不属于该数据集", img.Filename)
|
||
}
|
||
}
|
||
flag := 0
|
||
if req.Exclude {
|
||
flag = 1
|
||
}
|
||
if err := common.Serial().Submit(ctx, func() error {
|
||
return dao.DatasetImage.UpdateCleanExcluded(ctx, ids, flag)
|
||
}); err != nil {
|
||
return nil, err
|
||
}
|
||
excluded, err := dao.DatasetImage.CountExcludedByDataset(ctx, dataset.Id)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return &dto.AdminDatasetCleanApplyRes{Applied: len(ids), Excluded: excluded}, nil
|
||
}
|
||
|
||
// imageDHashFile 读图并算整图 64 位 dHash（读失败/解码失败返回 ok=false，调用方宁可保留）
|
||
func imageDHashFile(ctx context.Context, dir, filename string) (uint64, bool) {
|
||
data, err := os.ReadFile(filepath.Join(dir, filename))
|
||
if err != nil {
|
||
g.Log().Debugf(ctx, "数据清洗读图失败 %s: %+v", filename, err)
|
||
return 0, false
|
||
}
|
||
h, err := common.ImageDHash64(data)
|
||
if err != nil {
|
||
g.Log().Debugf(ctx, "数据清洗解码失败 %s: %+v", filename, err)
|
||
return 0, false
|
||
}
|
||
return h, true
|
||
}
|
||
|
||
// minHamming 与已保留哈希集的最小汉明距离（空集 = 无可比对象，返回 65 恒过阈值）
|
||
func minHamming(h uint64, hashes []uint64) int {
|
||
min := 65
|
||
for _, x := range hashes {
|
||
if d := bits.OnesCount64(h ^ x); d < min {
|
||
min = d
|
||
}
|
||
}
|
||
return min
|
||
}
|
||
|
||
// cleanBucketIdx 框高占比(%)落档（下含上不含；float 边界未命中回退最大档）
|
||
func cleanBucketIdx(pct float64) int {
|
||
for i, b := range consts.CleanBuckets {
|
||
if pct >= b.MinPct && pct < b.MaxPct {
|
||
return i
|
||
}
|
||
}
|
||
return len(consts.CleanBuckets) - 1
|
||
}
|