1471 lines
58 KiB
Go
1471 lines
58 KiB
Go
package service
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"image"
|
||
"image/jpeg"
|
||
"io"
|
||
"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"
|
||
"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 _, 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 {
|
||
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,
|
||
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)
|
||
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 模型文件随命名变更迁移（基名 = 文件名前缀，空回退数据集名）：
|
||
// 旧文件不存在或新旧路径相同（前缀未变且非空）直接跳过；失败仅记日志不阻断业务——
|
||
// 模型文件为付费训练产物，保留旧名总比删除好（下次训练直写新名覆盖）。
|
||
func migrateModelFile(ctx context.Context, oldName, oldPrefix, newName, newPrefix string) {
|
||
oldPath := common.TrainingModelPath(ctx, modelFileName(oldName, oldPrefix))
|
||
newPath := common.TrainingModelPath(ctx, modelFileName(newName, newPrefix))
|
||
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) {
|
||
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
|
||
}
|
||
|
||
// 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 生成任务占显存时跳过（同步等待不现实）
|
||
running, err := dao.GenTask.ListRunning(ctx)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
if len(running) > 0 {
|
||
return "", gerror.New("生成任务进行中（z-image 占用显存），未生成封面")
|
||
}
|
||
prompt := fmt.Sprintf("野外实拍照片：一只雄性%s和一只雌性%s并排站立在开阔的自然栖息地中，"+
|
||
"雄雌各一只清晰可见，16:9 横幅构图，真实照片质感，光线自然，画面清晰美观", species, species)
|
||
genCtx, cancel := context.WithTimeout(ctx, 120*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（调用方决定是否阻断）；
|
||
// 校验通过返回 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 // 未配置标注服务：不生成也不报错（池为空，生成图片时提示补参数）
|
||
}
|
||
// 显存互斥：z-image 生成任务占显存时跳过（同步等待不现实——生成任务可达小时级）
|
||
running, err := dao.GenTask.ListRunning(ctx)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if len(running) > 0 {
|
||
return nil, gerror.New("生成任务进行中（z-image 占用显存），暂无法生成参数，稍后重试")
|
||
}
|
||
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 输入失败")
|
||
}
|
||
genCtx, cancel := context.WithTimeout(ctx, 90*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 标注任务 / 该数据集训练进行中 / 已发布模型版本时拒绝
|
||
// （训练产物与模型为付费资产，需先删除模型版本再删数据集）。
|
||
// 删除 = 删图片/模型目录 + 删记录（标注随图片行删除，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)
|
||
}
|
||
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("该数据集有训练任务进行中，无法删除")
|
||
}
|
||
// 模型版本记录随数据集级联删除（管理端无模型管理界面，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, modelName)}
|
||
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
|
||
}
|
||
|
||
// singleAnimalClause 目标数量约束：每张仅一只动物，手填提示词统一追加（模板已由 {count} 固定 "1只"）
|
||
const singleAnimalClause = "画面中有且只有这一只动物，没有任何其他动物"
|
||
|
||
// buildPromptFromTemplate 按 config 通用模板组装提示词（单物种规则：物种固定取数据集 gen_species，
|
||
// 空则数据集名本身——生成表单无物种输入，不随机）。
|
||
// 物种只写名字不写羽毛细节（模型对正确名称自带外观先验）；场景/动作/遮挡从数据集表池随机
|
||
// （创建时 VLM 生成，config.yml 不再兜底；池空报错——生成图按池组装，池质量决定训练数据质量）；
|
||
// 光线走 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 := pickCfgList(poolOf(dataset.GenScenes))
|
||
action := pickCfgList(poolOf(dataset.GenActions))
|
||
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)
|
||
// 尺寸提示：按物理公式把表单距离换算成目标在图中的像素高与画面占比
|
||
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 供自动标注裁剪）
|
||
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,
|
||
} {
|
||
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))]
|
||
}
|
||
|
||
// 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 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, "创建图片目录失败")
|
||
}
|
||
// 命名：数据集属性 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",
|
||
Prompt: prompt,
|
||
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 != "" {
|
||
s.genTaskAutoLabel(bgCtx, dataset, addedIds)
|
||
// 部分失败：已入库图按付费资产保留，计数须计入（曾漏更新致 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
|
||
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)
|
||
}
|
||
s.genTaskAutoLabel(bgCtx, dataset, addedIds)
|
||
_ = dao.GenTask.Finish(bgCtx, taskId, "")
|
||
}()
|
||
}
|
||
|
||
// genTaskAutoLabel 生成任务完成后对本次新增图触发自动标注：忙（已有 running 任务）不报错，
|
||
// 由标注任务完成后的自动补标轮兜底；其他失败只记日志（付费资产，不可删）。
|
||
func (s *datasetService) genTaskAutoLabel(ctx context.Context, dataset *entity.Dataset, addedIds []int64) {
|
||
if len(addedIds) == 0 {
|
||
return
|
||
}
|
||
newImages, err := dao.DatasetImage.GetByIds(ctx, addedIds)
|
||
if err != nil {
|
||
g.Log().Errorf(ctx, "生成任务后读取新增图失败: %+v", err)
|
||
return
|
||
}
|
||
if err := LabelTask.AutoLabel(ctx, dataset, newImages); err != nil {
|
||
g.Log().Errorf(ctx, "生成任务后自动标注触发失败: %+v", err)
|
||
}
|
||
}
|
||
|
||
// 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,
|
||
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
|
||
}
|