【Sapphire开发日志 九】数据集查询接口封装

介绍

在完成用户加入和退出数据集的功能后，我们需要对数据集的查询结果进行封装，以便判断数据集是否为用户创建或认领的。本文将详细介绍如何实现这一功能。

具体实现

我们将通过封装数据集的查询结果，来判断用户是否为数据集的创建者或认领者。以下是具体的实现代码和详细解释：

type DatasetItem struct {
	ImgUrl       string `json:"imgUrl"`
	EmbeddingUrl string `json:"embeddingUrl"`
	Status       string `json:"status"`
	Id           int    `json:"id"`
}

type DatasetResult struct {
	DatasetId       uint          `json:"dataSetId"`
	DatasetName     string        `json:"dataSetName"`
	TaskInfo        string        `json:"taskInfo"`
	ObjectCnt       int           `json:"objectCnt"`
	Objects         []string      `json:"objects"`
	Owner           bool          `json:"owner"`
	Claim           bool          `json:"claim"`
	Status          string        `json:"status"`
	Datas           []DatasetItem `json:"datas"`
	Schedule        string        `json:"schedule"`
	TotalCount      int           `json:"totalCount"`
	EmbeddingCount  int           `json:"embeddingCount"`
	AnnotationCount int           `json:"annotationCount"`
	Finished        int           `json:"finished"`
}

func NewDatasetResult(dataset *domain.Dataset, isOwner bool, isClaim bool) *DatasetResult {
	var err error
	objects := make([]string, 0)
	if dataset.Tags != "" {
		tags := strings.Split(dataset.Tags, ",")
		for _, tag := range tags {
			if strings.TrimSpace(tag) == "" {
				continue
			} else {
				objects = append(objects, tag)
			}
		}
	}

	allImages, err := datasetDomain.ListImagesByStatusAndDatasetID(dataset.ID, -1)
	if err != nil {
		return nil
	}

	embeddingImages, err := datasetDomain.ListImagesByStatusAndDatasetID(dataset.ID, domain.ImgStatusEmbedded)
	if err != nil {
		return nil
	}

	annotationImages, err := datasetDomain.ListImagesByStatusAndDatasetID(dataset.ID, domain.ImgStatusAnnotated)
	if err != nil {
		return nil
	}

	// 将所有标注完成的图片加入到embeddingImages中
	embeddingImages = append(embeddingImages, annotationImages...)

	var statusStr string
	if len(allImages) == len(annotationImages) {
		statusStr = "annotationSuccess"
	} else if len(embeddingImages) == len(allImages) {
		statusStr = "Ready"
	} else {
		statusStr = "default"
	}

	return &DatasetResult{
		DatasetId:       dataset.ID,
		DatasetName:     dataset.Name,
		TaskInfo:        dataset.Description,
		ObjectCnt:       len(objects),
		Objects:         objects,
		Owner:           isOwner,
		Claim:           isClaim,
		Schedule:        dataset.EndTime.Format("2006-01-02 15:04:05"),
		TotalCount:      len(allImages),
		EmbeddingCount:  len(embeddingImages),
		AnnotationCount: len(annotationImages),
		Status:          statusStr,
		Finished:        0,
	}
}

func newDatasetItem(data *domain.ImgDataset) DatasetItem {
	var statusStr string
	switch data.Status {
	case domain.ImgStatusEmbedded:
		statusStr = "embedded"
	case domain.ImgStatusDefault:
		statusStr = "default"
	case domain.ImgStatusAnnotated:
		statusStr = "annotated"
	default:
		statusStr = "default"
	}

	return DatasetItem{
		ImgUrl:       data.ImgUrl,
		EmbeddingUrl: data.EmbeddingUrl,
		Status:       statusStr,
		Id:           int(data.ID),
	}
}

var datasetDomain = domain.NewDatasetDomain()

// GetAllDatasetList 获取数据集列表
func (s *DatasetService) GetAllDatasetList(userID uint) []*DatasetResult {
	var err error
	datasets, err := datasetDomain.GetDatasetList()
	if err != nil {
		// 返回空列表
		return make([]*DatasetResult, 0)
	}

	// 读取用户创建的数据集和加入的数据集
	userCreatedDatasets, err := datasetDomain.ListUserCreatedDatasets(userID)
	if err != nil {
		return make([]*DatasetResult, 0)
	}
	userCreatedMap := make(map[int]bool)
	for _, dataset := range userCreatedDatasets {
		userCreatedMap[int(dataset.ID)] = true
	}
	userJoinedDatasets, err := datasetDomain.ListUserJoinedDatasetList(userID)
	if err != nil {
		return make([]*DatasetResult, 0)
	}
	userJoinedMap := make(map[int]bool)
	for _, dataset := range userJoinedDatasets {
		userJoinedMap[int(dataset.ID)] = true
	}

	// 构建结果列表
	results := make([]*DatasetResult, 0)
	for _, dataset := range datasets {
		isOwner := dataset.CreatorID == userID
		isClaim := datasetDomain.IsUserClaimDataset(userID, dataset.ID)
		result := NewDatasetResult(&dataset, isOwner, isClaim)
		results = append(results, result)
	}
	return results
}

// GetUserCreatedDatasetList 获取用户创建的数据集列表
func (s *DatasetService) GetUserCreatedDatasetList(creatorID uint) []*DatasetResult {
	var err error
	datasets, err := datasetDomain.ListUserCreatedDatasets(creatorID)
	if err != nil {
		// 返回空列表
		return make([]*DatasetResult, 0)
	}

	results := s.buildResultList(datasets, true, false)

	return results
}

// GetUserJoinedDatasetList 获取用户加入的数据集列表
func (s *DatasetService) GetUserJoinedDatasetList(userID uint) []*DatasetResult {
	datasets, err := datasetDomain.ListUserJoinedDatasetList(userID)
	if err != nil {
		return make([]*DatasetResult, 0)
	}

	results := s.buildResultList(datasets, false, true)

	return results
}

// GetUserDatasetList 获取用户的数据集列表
func (s *DatasetService) GetUserDatasetList(userID uint) []*DatasetResult {
	createdDatasets, err := datasetDomain.ListUserCreatedDatasets(userID)
	if err != nil {
		return make([]*DatasetResult, 0)
	}
	createdResults := s.buildResultList(createdDatasets, true, false)

	joinedDatasets, err := datasetDomain.ListUserJoinedDatasetList(userID)
	if err != nil {
		return make([]*DatasetResult, 0)
	}
	joinedResults := s.buildResultList(joinedDatasets, false, true)

	// 创建一个新的列表，用于存储用户创建的数据集和加入的数据集
	// 同时创建一个map检查是否有重复的数据集
	results := make([]*DatasetResult, 0)
	datasetIdMap := make(map[uint]bool)
	allDatasets := append(createdResults, joinedResults...)
	for _, result := range allDatasets {
		// 当数据集ID已经存在时，跳过
		if _, ok := datasetIdMap[result.DatasetId]; ok {
			continue
		}
		results = append(results, result)
		datasetIdMap[result.DatasetId] = true
	}

	return results
}

代码解释

DatasetItem 和 DatasetResult 结构体

DatasetItem 和 DatasetResult 结构体用于封装数据集的详细信息和查询结果。

• DatasetItem 包含图片的URL、Embedding URL、状态和ID。
• DatasetResult 包含数据集的基本信息、任务信息、对象列表、用户是否是创建者、是否认领、数据集状态、数据项列表以及统计信息。

NewDatasetResult 函数

NewDatasetResult 函数用于创建一个新的 DatasetResult 对象。

1. 处理标签：将标签字符串拆分为标签数组。
2. 获取图片列表：根据图片状态获取所有图片、嵌入图片和标注图片。
3. 计算数据集状态：根据图片状态计算数据集的整体状态。
4. 封装结果：将所有信息封装到 DatasetResult 对象中并返回。

newDatasetItem 函数

newDatasetItem 函数用于创建一个新的 DatasetItem 对象。

1. 根据状态设置状态字符串：根据图片的状态设置对应的状态字符串。
2. 封装结果：将图片信息封装到 DatasetItem 对象中并返回。

GetAllDatasetList 函数

GetAllDatasetList 函数用于获取所有数据集的列表，并判断用户是否为数据集的创建者或认领者。

1. 获取所有数据集：从数据库中获取所有数据集。
2. 获取用户创建和加入的数据集：分别获取用户创建和加入的数据集，并存储到map中。
3. 构建结果列表：遍历所有数据集，判断用户是否为创建者或认领者，并封装结果。

GetUserCreatedDatasetList 函数

GetUserCreatedDatasetList 函数用于获取用户创建的数据集列表。

1. 获取用户创建的数据集：从数据库中获取用户创建的数据集。
2. 构建结果列表：将数据集信息封装到结果列表中。

GetUserJoinedDatasetList 函数

GetUserJoinedDatasetList 函数用于获取用户加入的数据集列表。

1. 获取用户加入的数据集：从数据库中获取用户加入的数据集。
2. 构建结果列表：将数据集信息封装到结果列表中。

GetUserDatasetList 函数

GetUserDatasetList 函数用于获取用户的数据集列表，包括用户创建和加入的数据集。

1. 获取用户创建和加入的数据集：分别获取用户创建和加入的数据集。
2. 构建结果列表：将创建和加入的数据集合并到一个列表中，并去重。

总结

本文详细介绍了如何封装数据集的查询结果，以判断用户是否为数据集的创建者或认领者。通过封装 DatasetItem 和 DatasetResult 结构体，并实现相关的查询函数，我们可以方便地获取用户相关的数据集信息，并进行相应的处理。