hpds_jkw_web/internal/service/dataset.go

package service

import (
	"context"
	"encoding/base64"
	"fmt"
	"git.hpds.cc/Component/logging"
	"hpds-iot-web/config"
	"hpds-iot-web/internal/proto"
	"hpds-iot-web/model"
	"hpds-iot-web/pkg/utils"
	"math"
	"math/rand"
	"net/http"
	"path"
	"time"
	"xorm.io/xorm"
)

type DatasetService interface {
	GetOwnerProjectList(ctx context.Context, req proto.OwnerProjectRequest) (rsp *proto.BaseResponse, err error)
	DatasetList(ctx context.Context, req proto.DatasetRequest) (rsp *proto.BaseResponse, err error)
	ImportDataset(ctx context.Context, req proto.ImportDatasetRequest) (rsp *proto.BaseResponse, err error)
	DatasetInfo(ctx context.Context, req proto.DatasetItemRequest) (rsp *proto.BaseResponse, err error)

	CreateTrainDataset(ctx context.Context, req proto.TrainDatasetRequest) (rsp *proto.BaseResponse, err error)
	TrainDatasetList(ctx context.Context, req proto.TrainDatasetItemRequest) (rsp *proto.BaseResponse, err error)
	TrainDatasetFileList(ctx context.Context, req proto.TrainDatasetItemRequest) (rsp *proto.BaseResponse, err error)
}

func NewDatasetService(cfg *config.WebConfig, engine *xorm.Engine, logger *logging.Logger) DatasetService {
	return &repo{
		AppConfig: cfg,
		engine:    engine,
		logger:    logger,
	}
}

func (rp *repo) GetOwnerProjectList(ctx context.Context, req proto.OwnerProjectRequest) (rsp *proto.BaseResponse, err error) {
	rsp = new(proto.BaseResponse)
	select {
	case <-ctx.Done():
		err = fmt.Errorf("超时/取消")
		rsp.Code = http.StatusInternalServerError
		rsp.Status = http.StatusText(http.StatusInternalServerError)
		rsp.Message = "超时/取消"
		rsp.Err = ctx.Err()
		return rsp, ctx.Err()
	default:
		ownerList := make([]model.Owner, 0)
		err := rp.engine.Where("(? = '' or owner_name like ?)", req.Key, "%"+req.Key+"%").
			And("status = 1").Find(&ownerList)
		if err != nil {
			goto ReturnPoint
		}
		data := make([]proto.OwnerProjectItem, len(ownerList))
		for k, v := range ownerList {
			projectList := make([]proto.ProjectItem, 0)
			err = rp.engine.Table("project").Cols("concat('"+fmt.Sprintf("%d", v.OwnerId)+"-', project_id) as project_id", "project_name").
				Where("owner_id = ?", v.OwnerId).And("status = 1").Find(&projectList)
			if err != nil {
				goto ReturnPoint
			}
			data[k] = proto.OwnerProjectItem{
				OwnerId:     v.OwnerId,
				OwnerName:   fmt.Sprintf("%s[%s]", v.OwnerName, v.ChargeUser),
				ProjectList: projectList,
			}
		}
		rsp.Code = http.StatusOK
		rsp.Status = http.StatusText(http.StatusOK)
		rsp.Message = "成功"
		rsp.Data = data
		rsp.Err = err
		return rsp, err
	}
ReturnPoint:
	if err != nil {
		rsp.Code = http.StatusInternalServerError
		rsp.Status = http.StatusText(http.StatusInternalServerError)
		rsp.Err = err
		rsp.Message = "失败"
	}
	return rsp, err
}

func (rp *repo) DatasetList(ctx context.Context, req proto.DatasetRequest) (rsp *proto.BaseResponse, err error) {
	rsp = new(proto.BaseResponse)
	select {
	case <-ctx.Done():
		err = fmt.Errorf("超时/取消")
		rsp.Code = http.StatusInternalServerError
		rsp.Status = http.StatusText(http.StatusInternalServerError)
		rsp.Message = "超时/取消"
		rsp.Err = ctx.Err()
		return rsp, ctx.Err()
	default:
		list := make([]model.Dataset, 0)
		st := rp.engine.Where("(? = '' or dataset_name like ?)", req.DatasetName, "%"+req.DatasetName+"%").
			And("(? = '' or create_at >= ?)", req.StartTime, req.StartTime).
			And("(? = '' or create_at < ?)", req.EndTime, req.EndTime).
			And("(? = 0 or owner_id = ?)", req.OwnerId, req.OwnerId).
			And("status = 1")
		if len(req.ProjectId) > 0 {
			st.In("project_id", req.ProjectId)
		}
		count, err := st.Limit(int(req.Size), int(((req.Page)-1)*req.Size)).FindAndCount(&list)
		if err != nil {
			goto ReturnPoint
		}
		data := make([]proto.DatasetItem, len(list))
		for k, v := range list {
			detailList := make([]model.FileManager, 0)
			datasetCount, err := rp.engine.Where("dataset_id = ?", v.DatasetId).And("data_type=1").FindAndCount(&detailList)
			if err != nil {
				goto ReturnPoint
			}
			fm := new(model.FileManager)
			datasetSize, err := rp.engine.Where("dataset_id = ?", v.DatasetId).And("data_type=1").SumInt(fm, "file_size")
			if err != nil {
				goto ReturnPoint
			}
			data[k] = proto.DatasetItem{
				DatasetId:    v.DatasetId,
				DatasetName:  v.DatasetName,
				DatasetDesc:  v.DatasetDesc,
				StoreName:    v.StoreName,
				CategoryId:   v.CategoryId,
				ProjectId:    v.ProjectId,
				OwnerId:      v.OwnerId,
				Creator:      v.Creator,
				CreateAt:     v.CreateAt,
				DatasetCount: datasetCount,
				DatasetSize:  datasetSize,
			}
		}
		rsp.Code = http.StatusOK
		rsp.Status = http.StatusText(http.StatusOK)
		rsp.Message = "成功"
		rsp = FillPaging(count, req.Page, req.Size, data, rsp)
		rsp.Err = err
		return rsp, err
	}
ReturnPoint:
	if err != nil {
		rsp.Code = http.StatusInternalServerError
		rsp.Status = http.StatusText(http.StatusInternalServerError)
		rsp.Err = err
		rsp.Message = "失败"
	}
	return rsp, err
}

func (rp *repo) ImportDataset(ctx context.Context, req proto.ImportDatasetRequest) (rsp *proto.BaseResponse, err error) {
	rsp = new(proto.BaseResponse)
	select {
	case <-ctx.Done():
		err = fmt.Errorf("超时/取消")
		rsp.Code = http.StatusInternalServerError
		rsp.Status = http.StatusText(http.StatusInternalServerError)
		rsp.Message = "超时/取消"
		rsp.Err = ctx.Err()
		return rsp, ctx.Err()
	default:
		item := &model.Dataset{
			DatasetName: req.DatasetName,
			DatasetDesc: req.DatasetDesc,
			CategoryId:  req.CategoryId,
			ProjectId:   req.ProjectId,
			OwnerId:     req.OwnerId,
			StoreName:   req.StoreName,
			Creator:     req.Creator,
			Status:      1,
			CreateAt:    time.Now().Unix(),
			UpdateAt:    time.Now().Unix(),
		}
		_, err = rp.engine.Insert(item)
		if err != nil {
			goto ReturnPoint
		}
		rsp.Code = http.StatusOK
		rsp.Status = http.StatusText(http.StatusOK)
		rsp.Message = "成功"
		rsp.Data = item
		rsp.Err = err
		return rsp, err
	}
ReturnPoint:
	if err != nil {
		rsp.Code = http.StatusInternalServerError
		rsp.Status = http.StatusText(http.StatusInternalServerError)
		rsp.Err = err
		rsp.Message = "失败"
	}
	return rsp, err
}

func (rp *repo) DatasetInfo(ctx context.Context, req proto.DatasetItemRequest) (rsp *proto.BaseResponse, err error) {
	rsp = new(proto.BaseResponse)
	select {
	case <-ctx.Done():
		err = fmt.Errorf("超时/取消")
		rsp.Code = http.StatusInternalServerError
		rsp.Status = http.StatusText(http.StatusInternalServerError)
		rsp.Message = "超时/取消"
		rsp.Err = ctx.Err()
		return rsp, ctx.Err()
	default:
		item := new(model.Dataset)
		var b bool
		b, err = rp.engine.ID(req.DatasetId).Get(item)
		if err != nil {
			goto ReturnPoint
		}
		if !b {
			err = fmt.Errorf("未能找到对应的数据集")
			goto ReturnPoint
		}

		rsp.Code = http.StatusOK
		rsp.Status = http.StatusText(http.StatusOK)
		rsp.Message = "成功"
		rsp.Err = err
		rsp.Data = item
		return rsp, err
	}
ReturnPoint:
	if err != nil {
		rsp.Code = http.StatusInternalServerError
		rsp.Status = http.StatusText(http.StatusInternalServerError)
		rsp.Err = err
		rsp.Message = "失败"
	}
	return rsp, err
}

func (rp *repo) CreateTrainDataset(ctx context.Context, req proto.TrainDatasetRequest) (rsp *proto.BaseResponse, err error) {
	rsp = new(proto.BaseResponse)
	select {
	case <-ctx.Done():
		err = fmt.Errorf("超时/取消")
		rsp.Code = http.StatusInternalServerError
		rsp.Status = http.StatusText(http.StatusInternalServerError)
		rsp.Message = "超时/取消"
		rsp.Err = ctx.Err()
		return rsp, ctx.Err()
	default:
		var (
			h             bool
			trainFileList []model.FileManager
			valFileList   []model.FileManager
			testFileList  []model.FileManager
			//wg            sync.WaitGroup
		)
		log := new(model.DatasetOperationLog)
		dataset := new(model.Dataset)
		h, err = rp.engine.ID(req.DatasetId).Get(dataset)
		if err != nil {
			goto ReturnPoint
		}
		if !h {
			err = fmt.Errorf("未能找到对应的采集数据集")
			goto ReturnPoint
		}
		fileList := make([]model.FileManager, 0)
		err = rp.engine.Where("dataset_id = ?", req.DatasetId).And("is_disease > 0").Find(&fileList)
		if err != nil {
			goto ReturnPoint
		}
		if req.TargetData == 0 {
			req.TargetData = len(fileList)
		}
		if req.TargetData > len(fileList) {
			err = fmt.Errorf("超出现有标注数据集数量")
			goto ReturnPoint
		}
		if req.SplitMethod == 1 { //随机
			rand.Seed(time.Now().UnixNano())
			rand.Shuffle(len(fileList), func(i, j int) {
				fileList[i], fileList[j] = fileList[j], fileList[i]
			})
		}
		trainNumber := int(math.Floor(float64(int64(req.TargetData)*req.TrainNumber) / 100))
		valNumber := int(math.Floor(float64(int64(req.TargetData)*req.ValidationNumber) / 100))
		testNumber := req.TargetData - trainNumber - valNumber
		if trainNumber-1 > 1 {
			trainFileList = fileList[:trainNumber-1]
		} else {
			trainFileList = make([]model.FileManager, 0)
			trainFileList = append(trainFileList, fileList[0])
		}
		if trainNumber != trainNumber+valNumber-1 {
			valFileList = fileList[trainNumber : trainNumber+valNumber-1]
		} else {
			valFileList = make([]model.FileManager, 0)
			valFileList = append(valFileList, fileList[trainNumber])
		}
		if trainNumber+valNumber < len(fileList) {
			testFileList = fileList[trainNumber+valNumber:]
		} else {
			testFileList = make([]model.FileManager, 0)
			testFileList = append(testFileList, fileList[trainNumber+valNumber])
		}

		train := new(model.TrainingDataset)
		h, err = rp.engine.Where("name = ?", req.TrainName).Get(train)
		if err != nil {
			goto ReturnPoint
		}
		if !h {
			train.Name = req.TrainName
			train.DatasetDesc = req.TrainDesc
			//train.DatasetId = req.DatasetId
			train.CategoryId = dataset.CategoryId
			_, err = rp.engine.Insert(train)
			if err != nil {
				goto ReturnPoint
			}
		}
		log.TargetData = int64(req.TargetData)
		log.DatasetId = req.DatasetId
		log.TrainingDatasetId = train.DatasetId
		log.SplitMethod = req.SplitMethod
		log.TrainNumber = int64(trainNumber)
		log.ValidationNumber = int64(valNumber)
		log.TestNumber = int64(testNumber)
		log.Creator = req.UserId
		_, err = rp.engine.Insert(log)
		if err != nil {
			goto ReturnPoint
		}
		//wg.Add(3)
		go BatchCopyData(trainFileList, train.DatasetId, log.LogId, req.UserId, 1, req.TrainName, rp) //, &wg
		go BatchCopyData(valFileList, train.DatasetId, log.LogId, req.UserId, 2, req.TrainName, rp)   //, &wg
		go BatchCopyData(testFileList, train.DatasetId, log.LogId, req.UserId, 3, req.TrainName, rp)  //, &wg
		//wg.Wait()
		rsp.Code = http.StatusOK
		rsp.Status = http.StatusText(http.StatusOK)
		rsp.Message = "成功"
		rsp.Err = err
		rsp.Data = log
		return rsp, err
	}
ReturnPoint:
	if err != nil {
		rsp.Code = http.StatusInternalServerError
		rsp.Status = http.StatusText(http.StatusInternalServerError)
		rsp.Err = err
		rsp.Message = "失败"
	}
	return rsp, err
}

func BatchCopyData(list []model.FileManager, trainId, logId, userId int64, categoryId int, trainName string, rp *repo) { //, wg *sync.WaitGroup
	batchList := make([]model.TrainingDatasetDetail, len(list))
	for k, v := range list {
		dir := "no_disease"
		if v.IsDisease == 1 {
			dir = "disease"
		}
		utils.DownloadMinioFileToLocalPath(v.AccessUrl, path.Join(rp.AppConfig.TrainDir, trainName, model.GetTrainCategory(categoryId), dir), v.FileName,
			rp.AppConfig.Minio.Protocol, rp.AppConfig.Minio.Endpoint, rp.AppConfig.Minio.Bucket, rp.AppConfig.Minio.AccessKeyId,
			rp.AppConfig.Minio.SecretAccessKey, rp.logger)
		item := model.TrainingDatasetDetail{
			FileName:       v.FileName,
			FilePath:       path.Join(rp.AppConfig.TrainDir, trainName, model.GetTrainCategory(categoryId), dir, v.FileName),
			DatasetId:      trainId,
			CategoryId:     categoryId,
			FileSize:       v.FileSize,
			FileMd5:        v.FileMd5,
			IsDisease:      v.IsDisease,
			OperationLogId: logId,
			Creator:        userId,
			CreateAt:       time.Now().Unix(),
			UpdateAt:       time.Now().Unix(),
		}
		batchList[k] = item
	}
	_, _ = rp.engine.Insert(batchList)
	//wg.Done()
}

func (rp *repo) TrainDatasetList(ctx context.Context, req proto.TrainDatasetItemRequest) (rsp *proto.BaseResponse, err error) {
	rsp = new(proto.BaseResponse)
	select {
	case <-ctx.Done():
		err = fmt.Errorf("超时/取消")
		rsp.Code = http.StatusInternalServerError
		rsp.Status = http.StatusText(http.StatusInternalServerError)
		rsp.Message = "超时/取消"
		rsp.Err = ctx.Err()
		return rsp, ctx.Err()
	default:
		type QuantityStatistics struct {
			DatasetId  int64
			CategoryId int
			Total      int64
		}

		var (
			count int64
			list  []proto.TrainingDataset
		)
		trainingList := make([]model.TrainingDataset, 0)
		count, err = rp.engine.Where("(?=0 or dataset_id = ?)", req.DatasetId, req.DatasetId).
			And("(?= 0 or category_id = ?)", req.BizType, req.BizType).
			And("(? ='' or name like ?)", req.TrainName, "%"+req.TrainName+"%").
			Limit(int(req.Size), int(((req.Page)-1)*req.Size)).FindAndCount(&trainingList)
		if err != nil {
			goto ReturnPoint
		}
		list = make([]proto.TrainingDataset, len(trainingList))
		for k, v := range trainingList {
			qs := make([]QuantityStatistics, 0)
			err = rp.engine.SQL("select dataset_id, category_id, count(1) as total from training_dataset_detail where dataset_id = ? group by category_id, dataset_id", v.DatasetId).Find(&qs)
			if err != nil {
				goto ReturnPoint
			}
			item := proto.TrainingDataset{
				DatasetId:   v.DatasetId,
				Name:        v.Name,
				CategoryId:  v.CategoryId,
				DatasetDesc: v.DatasetDesc,
				TotalSize:   0,
				TrainSize:   0,
				ValSize:     0,
				TestSize:    0,
				StoreName:   v.StoreName,
				CreateAt:    v.CreateAt,
				UpdateAt:    v.UpdateAt,
			}
			for _, val := range qs {
				switch val.CategoryId {
				case 1:
					item.TrainSize = val.Total
				case 2:
					item.ValSize = val.Total
				case 3:
					item.TestSize = val.Total
				}
				item.TotalSize += val.Total
			}
			list[k] = item
		}
		rsp.Code = http.StatusOK
		rsp.Status = http.StatusText(http.StatusOK)
		rsp.Message = "成功"
		rsp = FillPaging(count, req.Page, req.Size, list, rsp)
		rsp.Err = err
		return rsp, err
	}

ReturnPoint:
	if err != nil {
		rsp.Code = http.StatusInternalServerError
		rsp.Status = http.StatusText(http.StatusInternalServerError)
		rsp.Err = err
		rsp.Message = "失败"
	}
	return rsp, err
}

func (rp *repo) TrainDatasetFileList(ctx context.Context, req proto.TrainDatasetItemRequest) (rsp *proto.BaseResponse, err error) {
	rsp = new(proto.BaseResponse)
	select {
	case <-ctx.Done():
		err = fmt.Errorf("超时/取消")
		rsp.Code = http.StatusInternalServerError
		rsp.Status = http.StatusText(http.StatusInternalServerError)
		rsp.Message = "超时/取消"
		rsp.Err = ctx.Err()
		return rsp, ctx.Err()
	default:
		var (
			count int64
			list  []proto.TrainingDatasetFileItem
		)
		fileList := make([]model.TrainingDatasetDetail, 0)
		count, err = rp.engine.Where("dataset_id = ?", req.DatasetId).
			Limit(int(req.Size), int(((req.Page)-1)*req.Size)).FindAndCount(&fileList)
		if err != nil {
			goto ReturnPoint
		}
		list = make([]proto.TrainingDatasetFileItem, len(fileList))
		for k, v := range fileList {
			buff := utils.ReadFile(v.FilePath)
			img := utils.BuffToImage(buff)
			buf := utils.ImageToBuff(img, "jpeg")
			list[k] = proto.TrainingDatasetFileItem{
				DetailId:    v.DetailId,
				FileName:    v.FileName,
				FileSize:    v.FileSize,
				FilePath:    v.FilePath,
				FileContent: "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString(buf.Bytes()),
				IsDisease:   v.IsDisease,
				CategoryId:  v.CategoryId,
			}
		}

		rsp.Code = http.StatusOK
		rsp.Status = http.StatusText(http.StatusOK)
		rsp.Message = "成功"
		rsp = FillPaging(count, req.Page, req.Size, list, rsp)
		rsp.Err = err
		return rsp, err
	}

ReturnPoint:
	if err != nil {
		rsp.Code = http.StatusInternalServerError
		rsp.Status = http.StatusText(http.StatusInternalServerError)
		rsp.Err = err
		rsp.Message = "失败"
	}
	return rsp, err
}