322 lines
12 KiB
Go
322 lines
12 KiB
Go
package dao
|
||
|
||
import (
|
||
"context"
|
||
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
"github.com/gogf/gf/v2/util/gconv"
|
||
|
||
"observer-server/biz/consts"
|
||
"observer-server/biz/model/entity"
|
||
"observer-server/common"
|
||
)
|
||
|
||
// DatasetImage 数据集图片表 DAO：文件名唯一（防重名），按数据集查询；表小不做缓存。
|
||
type datasetImageDao struct{}
|
||
|
||
var DatasetImage = &datasetImageDao{}
|
||
|
||
func init() {
|
||
ctx := context.Background()
|
||
_, err := g.DB().Exec(ctx, `CREATE TABLE IF NOT EXISTS dataset_image (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
dataset_id INTEGER NOT NULL,
|
||
filename TEXT NOT NULL,
|
||
source TEXT NOT NULL DEFAULT 'manual',
|
||
animal_count INTEGER NOT NULL DEFAULT 0,
|
||
labels_json TEXT,
|
||
created_at TEXT NOT NULL,
|
||
UNIQUE (dataset_id, filename)
|
||
)`)
|
||
if err != nil {
|
||
panic(err)
|
||
}
|
||
// 存量库迁移：标注列缺失时追加（新库建表已含列，幂等跳过）；candidates_json 已随 v10 删除
|
||
// 注意 EnsureColumn 的 ddl 参数须自带列名（拼接为 ALTER TABLE ADD COLUMN <ddl>）
|
||
common.EnsureColumn(ctx, consts.TableDatasetImage, "labels_json", "labels_json TEXT")
|
||
common.EnsureColumn(ctx, consts.TableDatasetImage, "animal_count", "animal_count INTEGER NOT NULL DEFAULT 0")
|
||
common.EnsureColumn(ctx, consts.TableDatasetImage, "clean_excluded", "clean_excluded INTEGER NOT NULL DEFAULT 0")
|
||
common.EnsureColumn(ctx, consts.TableDatasetImage, "review_status", "review_status INTEGER NOT NULL DEFAULT 0")
|
||
// 存量库迁移：众包任务图片粒度下发占用标记（2026-09-07，0=未下发）
|
||
common.EnsureColumn(ctx, consts.TableDatasetImage, "annotate_task_id", "annotate_task_id INTEGER NOT NULL DEFAULT 0")
|
||
// 存量迁移（幂等）：加列后既有标注统一视为已审定稿（非空框→2）；上线后不变式
|
||
// review_status=0 ⟹ labels_json 无框 恒成立（拒绝清框、提交/预标置 1、审核置 2），
|
||
// 故本 UPDATE 只会命中存量行，不会误改新数据
|
||
if _, err := g.DB().Exec(ctx, `UPDATE dataset_image SET review_status = 2
|
||
WHERE review_status = 0 AND labels_json IS NOT NULL AND labels_json != '' AND labels_json != '[]'`); err != nil {
|
||
panic(err)
|
||
}
|
||
// 存量库迁移：删除 prompt 列（生成提示词不再存储，2026-09-02 决策）
|
||
cols, err := g.DB().GetAll(ctx, "PRAGMA table_info(dataset_image)")
|
||
if err != nil {
|
||
panic(err)
|
||
}
|
||
for _, col := range cols {
|
||
if gconv.String(col["name"]) == "prompt" {
|
||
if _, err := g.DB().Exec(ctx, "ALTER TABLE dataset_image DROP COLUMN prompt"); err != nil {
|
||
panic(err)
|
||
}
|
||
g.Log().Warningf(ctx, "存量表 dataset_image 已迁移：删除 prompt 列")
|
||
break
|
||
}
|
||
}
|
||
}
|
||
|
||
// Insert 插入图片记录
|
||
func (d *datasetImageDao) Insert(ctx context.Context, m *entity.DatasetImage) (int64, error) {
|
||
res, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).Data(g.Map{
|
||
"dataset_id": m.DatasetId,
|
||
"filename": m.Filename,
|
||
"source": m.Source,
|
||
"animal_count": m.AnimalCount,
|
||
"created_at": m.CreatedAt,
|
||
}).Insert()
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
return res.LastInsertId()
|
||
}
|
||
|
||
// GetById 按主键取图片记录
|
||
func (d *datasetImageDao) GetById(ctx context.Context, id int64) (*entity.DatasetImage, error) {
|
||
var one *entity.DatasetImage
|
||
err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).Where("id", id).Scan(&one)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return one, nil
|
||
}
|
||
|
||
// ListByDataset 某数据集全部图片（按文件名编号倒序，详情页逐行展示用）。
|
||
// 文件名形如 <前缀>_<编号>.jpg（无前缀为时间戳数字名），须提取下划线后的数字按数值倒序：
|
||
// 字符串序会把 pheasant_101.jpg 排在 pheasant_2.jpg 之前（'1'<'9'），导致最大编号落到底部
|
||
func (d *datasetImageDao) ListByDataset(ctx context.Context, datasetId int64) ([]*entity.DatasetImage, error) {
|
||
var list []*entity.DatasetImage
|
||
err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).
|
||
Where("dataset_id", datasetId).
|
||
Order("CAST(substr(filename, instr(filename, '_') + 1) AS INTEGER) DESC").Scan(&list)
|
||
if err != nil {
|
||
if common.IsNoRows(err) {
|
||
return []*entity.DatasetImage{}, nil
|
||
}
|
||
return nil, err
|
||
}
|
||
return list, nil
|
||
}
|
||
|
||
// ListPoolByTask 众包任务池（技术设计.md「App 标注众包与时长激励」；2026-09-07 图片粒度下发）：
|
||
// 本任务占用（annotate_task_id=taskId）且未审核未清洗排除的图；待审核(1)/已审核(2)不进池——
|
||
// 分别等管理端审核与已定稿
|
||
func (d *datasetImageDao) ListPoolByTask(ctx context.Context, taskId int64) ([]*entity.DatasetImage, error) {
|
||
var list []*entity.DatasetImage
|
||
err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).
|
||
Where("annotate_task_id", taskId).
|
||
Where("review_status", consts.ReviewImageNone).
|
||
Where("clean_excluded", 0).
|
||
OrderAsc("id").Scan(&list)
|
||
if err != nil {
|
||
if common.IsNoRows(err) {
|
||
return []*entity.DatasetImage{}, nil
|
||
}
|
||
return nil, err
|
||
}
|
||
return list, nil
|
||
}
|
||
|
||
// CountPoolByTaskIds 多任务池余量（一次 GROUP BY；众包任务列表展示用）
|
||
func (d *datasetImageDao) CountPoolByTaskIds(ctx context.Context, taskIds []int64) (map[int64]int64, error) {
|
||
out := make(map[int64]int64)
|
||
if len(taskIds) == 0 {
|
||
return out, nil
|
||
}
|
||
var rows []struct {
|
||
AnnotateTaskId int64 `orm:"annotate_task_id"`
|
||
Cnt int64 `orm:"cnt"`
|
||
}
|
||
err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).
|
||
WhereIn("annotate_task_id", taskIds).
|
||
Where("review_status", consts.ReviewImageNone).
|
||
Where("clean_excluded", 0).
|
||
Fields("annotate_task_id, COUNT(*) AS cnt").
|
||
Group("annotate_task_id").
|
||
Scan(&rows)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
for _, r := range rows {
|
||
out[r.AnnotateTaskId] = r.Cnt
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// OccupyByTask 下发占用：批量写 annotate_task_id（Serial 内调用，防与并发下发竞态；IN ≤100 分批）
|
||
func (d *datasetImageDao) OccupyByTask(ctx context.Context, ids []int64, taskId int64) error {
|
||
for start := 0; start < len(ids); start += 100 {
|
||
end := start + 100
|
||
if end > len(ids) {
|
||
end = len(ids)
|
||
}
|
||
if _, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).
|
||
WhereIn("id", ids[start:end]).
|
||
Where("annotate_task_id", 0).
|
||
Data(g.Map{"annotate_task_id": taskId}).Update(); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// ReleaseByTask 停用释放：清指定图的任务占用标记回未标注（仅当仍占用本任务，防误清新任务标记；IN ≤100 分批）
|
||
func (d *datasetImageDao) ReleaseByTask(ctx context.Context, ids []int64, taskId int64) error {
|
||
for start := 0; start < len(ids); start += 100 {
|
||
end := start + 100
|
||
if end > len(ids) {
|
||
end = len(ids)
|
||
}
|
||
if _, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).
|
||
WhereIn("id", ids[start:end]).
|
||
Where("annotate_task_id", taskId).
|
||
Data(g.Map{"annotate_task_id": 0}).Update(); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// GetByIds 按 id 批量取（删除图片定位文件用）
|
||
func (d *datasetImageDao) GetByIds(ctx context.Context, ids []int64) ([]*entity.DatasetImage, error) {
|
||
var list []*entity.DatasetImage
|
||
err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).WhereIn("id", ids).Scan(&list)
|
||
if err != nil {
|
||
if common.IsNoRows(err) {
|
||
return []*entity.DatasetImage{}, nil
|
||
}
|
||
return nil, err
|
||
}
|
||
return list, nil
|
||
}
|
||
|
||
// GetByFilename 按文件名取（图片访问校验归属）
|
||
func (d *datasetImageDao) GetByFilename(ctx context.Context, datasetId int64, filename string) (*entity.DatasetImage, error) {
|
||
var e entity.DatasetImage
|
||
err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).
|
||
Where("dataset_id", datasetId).Where("filename", filename).Scan(&e)
|
||
if err != nil {
|
||
if common.IsNoRows(err) {
|
||
return nil, nil
|
||
}
|
||
return nil, err
|
||
}
|
||
return &e, nil
|
||
}
|
||
|
||
// CountByDataset 某数据集图片数
|
||
func (d *datasetImageDao) CountByDataset(ctx context.Context, datasetId int64) (int64, error) {
|
||
n, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).Where("dataset_id", datasetId).Count()
|
||
return int64(n), err
|
||
}
|
||
|
||
// UpdateLabels 覆写标注 JSON（整图粒度；空串=清空标注；AI 自动标注与人工保存共用）
|
||
func (d *datasetImageDao) UpdateLabels(ctx context.Context, id int64, labelsJson string) error {
|
||
_, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).Where("id", id).
|
||
Data(g.Map{"labels_json": labelsJson}).Update()
|
||
return err
|
||
}
|
||
|
||
// UpdateLabelsAndReview 覆写标注 JSON 并同步审核状态（预标/AdminLabelSave/App 提交共用）：
|
||
// 预标检出/人工保存/App 提交 → 1；人工清空保存 → 0（回到未标注池）
|
||
func (d *datasetImageDao) UpdateLabelsAndReview(ctx context.Context, id int64, labelsJson string, reviewStatus int) error {
|
||
_, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).Where("id", id).
|
||
Data(g.Map{"labels_json": labelsJson, "review_status": reviewStatus}).Update()
|
||
return err
|
||
}
|
||
|
||
// SetReviewStatus 只改审核状态（审核通过=2 / 拒绝先清标注再置 0，两步由 service 编排）
|
||
func (d *datasetImageDao) SetReviewStatus(ctx context.Context, ids []int64, reviewStatus int) error {
|
||
if len(ids) == 0 {
|
||
return nil
|
||
}
|
||
_, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).WhereIn("id", ids).
|
||
Data(g.Map{"review_status": reviewStatus}).Update()
|
||
return err
|
||
}
|
||
|
||
// CountLabeledByDataset 某数据集已审核图片数（review_status=2；2026-09-04 起已标注口径
|
||
// = 人工审核通过，labels_json 非空只是必要条件）
|
||
func (d *datasetImageDao) CountLabeledByDataset(ctx context.Context, datasetId int64) (int64, error) {
|
||
n, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).
|
||
Where("dataset_id", datasetId).
|
||
Where("review_status", consts.ReviewImageApproved).
|
||
Count()
|
||
return int64(n), err
|
||
}
|
||
|
||
// CountStat 某数据集图片数/已标注数（列表实时统计：冗余计数列会漂移，展示一律以此为准）
|
||
type CountStat struct {
|
||
Count int64
|
||
Labeled int64
|
||
}
|
||
|
||
// CountStatByDatasets 多数据集图片数/已标注数（一次 GROUP BY 查询；IN 按 ≤100 分批）
|
||
func (d *datasetImageDao) CountStatByDatasets(ctx context.Context, datasetIds []int64) (map[int64]CountStat, error) {
|
||
stat := make(map[int64]CountStat)
|
||
for start := 0; start < len(datasetIds); start += 100 {
|
||
end := start + 100
|
||
if end > len(datasetIds) {
|
||
end = len(datasetIds)
|
||
}
|
||
var rows []struct {
|
||
DatasetId int64 `orm:"dataset_id"`
|
||
Cnt int64 `orm:"cnt"`
|
||
Labeled int64 `orm:"labeled_cnt"`
|
||
}
|
||
err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).
|
||
WhereIn("dataset_id", datasetIds[start:end]).
|
||
Fields("dataset_id, COUNT(*) AS cnt, SUM(CASE WHEN review_status = 2 THEN 1 ELSE 0 END) AS labeled_cnt").
|
||
Group("dataset_id").
|
||
Scan(&rows)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
for _, r := range rows {
|
||
stat[r.DatasetId] = CountStat{Count: r.Cnt, Labeled: r.Labeled}
|
||
}
|
||
}
|
||
return stat, nil
|
||
}
|
||
|
||
// DeleteByIds 按 id 批量删除（标注随行删除）
|
||
func (d *datasetImageDao) DeleteByIds(ctx context.Context, ids []int64) error {
|
||
_, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).WhereIn("id", ids).Delete()
|
||
return err
|
||
}
|
||
|
||
// DeleteByDataset 删除某数据集全部图片记录（删数据集联动）
|
||
func (d *datasetImageDao) DeleteByDataset(ctx context.Context, datasetId int64) error {
|
||
_, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).Where("dataset_id", datasetId).Delete()
|
||
return err
|
||
}
|
||
|
||
// UpdateCleanExcluded 批量置/清数据清洗排除标记（excluded: 1=出训练集, 0=恢复;IN ≤100 分批）
|
||
func (d *datasetImageDao) UpdateCleanExcluded(ctx context.Context, ids []int64, excluded int) error {
|
||
for start := 0; start < len(ids); start += 100 {
|
||
end := start + 100
|
||
if end > len(ids) {
|
||
end = len(ids)
|
||
}
|
||
if _, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).
|
||
WhereIn("id", ids[start:end]).
|
||
Data(g.Map{"clean_excluded": excluded}).Update(); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// CountExcludedByDataset 某数据集已排除（clean_excluded=1）图数（数据清洗后角标/计数刷新）
|
||
func (d *datasetImageDao) CountExcludedByDataset(ctx context.Context, datasetId int64) (int64, error) {
|
||
n, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).
|
||
Where("dataset_id", datasetId).Where("clean_excluded", 1).Count()
|
||
return int64(n), err
|
||
}
|