4dkankan-center-openApi/src/main/java/com/fdkankan/openApi/service/www/impl/SceneMarkShapeServiceImpl.java

package com.fdkankan.openApi.service.www.impl;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.lang.UUID;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.baomidou.dynamic.datasource.annotation.DS;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.fdkankan.common.constant.ErrorCode;
import com.fdkankan.common.exception.BusinessException;
import com.fdkankan.model.constants.ConstantFilePath;
import com.fdkankan.openApi.dto.SceneJsonDTO;
import com.fdkankan.openApi.dto.SceneMarkShapesDTO;
import com.fdkankan.openApi.entity.www.SceneMarkShape;
import com.fdkankan.openApi.mapper.www.IMarkShapeMapper;
import com.fdkankan.openApi.service.www.ISceneMarkShapeService;
import com.fdkankan.openApi.util.ConverxyUtil;
import com.fdkankan.openApi.vo.www.SceneMarkShapeParamVO;
import com.fdkankan.openApi.vo.www.SceneMarkShapeReDetectParamVO;
import com.fdkankan.openApi.vo.www.SceneMarkShapeVO;
import com.fdkankan.rabbitmq.util.RabbitMqProducer;
import com.fdkankan.web.response.ResultData;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;


/**
 * Created by Xiewj on 2021/11/23 0026 10:14
 */
@Slf4j
@Service
@DS("www")
public class SceneMarkShapeServiceImpl extends ServiceImpl<IMarkShapeMapper, SceneMarkShape> implements ISceneMarkShapeService {
    @Autowired
    private RabbitMqProducer rabbitMqProducer;
    @Value("${queue.scene.yolov5-train-queue}")
    private String yolov5TrainQueue;
    @Value("${main.url}")
    private String mainUrl;
    @Autowired
    private ISceneMarkShapeService sceneMarkShapeService;

    @Override
    public ResultData saveOrEdit(SceneMarkShapeVO param) {
        if (Objects.isNull(param.getShapes())) {
            throw new BusinessException(ErrorCode.PARAM_REQUIRED,"shapes 为空");
        }
        log.info("标注数据保存或修改-场景码:{},场景图片:{},",param.getNum(),param.getImagePath());
        SceneMarkShape shape = this.findByNumAndImagePath(param.getNum(), param.getImagePath());
        SceneMarkShape sceneMarkShape=new SceneMarkShape();
        BeanUtil.copyProperties(param,sceneMarkShape);
        if (ObjectUtil.isNotNull(shape)){
            log.info("修改成功-修改id--{}",shape.getId());
            sceneMarkShape.setId(shape.getId());
            sceneMarkShape.setUpdateTime(new Date());
            return ResultData.ok(this.updateById(sceneMarkShape));
        }else {
            log.info("保存成功");
            sceneMarkShape.setCreateTime(new Date());
            return ResultData.ok(this.save(sceneMarkShape));
        }
    }

    @Override
    public ResultData saveOrEditBatch(SceneMarkShapesDTO param) {

        Set<String> imagePaths = param.getList().stream().map(v -> v.getImagePath()).collect(Collectors.toSet());
        List<SceneMarkShape> sceneMarkShapes = this.baseMapper.selectByNumAndImagePaths(param.getNum(), imagePaths);
        Map<String, Long> imagePathIdMap = sceneMarkShapes.stream().collect(Collectors.toMap(SceneMarkShape::getImagePath, SceneMarkShape::getId));

        List<SceneMarkShape> sceneMarkShapeList = BeanUtil.copyToList(param.getList(), SceneMarkShape.class);
        sceneMarkShapeList.stream().forEach(v -> v.setNum(param.getNum()));
        if(CollUtil.isNotEmpty(imagePathIdMap)){
            sceneMarkShapeList.stream().forEach(v->{
                v.setId(imagePathIdMap.get(v.getImagePath()));
            });
        }
        return ResultData.ok(this.saveOrUpdateBatch(sceneMarkShapeList));
    }


    @Override
    public ResultData getShapesInfo(SceneMarkShapeParamVO param) {
        SceneMarkShape res=  findByNumAndImagePath(param.getNum(),param.getImagePath());
        SceneMarkShapeVO vo=new SceneMarkShapeVO();
        BeanUtil.copyProperties(res,vo);
        return ResultData.ok(vo);
    }

    @Override
    public void editTrainStatus(SceneMarkShapeParamVO param) {
        SceneMarkShape byNumAndImagePath = findByNumAndImagePath(param.getNum(), param.getImagePath());
        if (ObjectUtil.isNotNull(byNumAndImagePath)){
            byNumAndImagePath.setReTrain(0);
            byNumAndImagePath.setToTrain(1);
            updateById(byNumAndImagePath);
        }
    }

    @Override
    public List<SceneMarkShape> findByReTrainStatus(Integer reTrain){
        LambdaQueryWrapper<SceneMarkShape> wrapper = Wrappers.lambdaQuery();
        wrapper.eq(SceneMarkShape::getReTrain,reTrain);
        return list(wrapper);
    }

    @Override
    public List<SceneMarkShape> findByToTrainStatus(Integer toTrain) {
        QueryWrapper queryWrapper = new QueryWrapper();
        queryWrapper.select("DISTINCT num")
                .eq("to_train",toTrain) ;
        return getBaseMapper().selectList(queryWrapper);
    }

    @Override
    public void trainScene(SceneMarkShapeReDetectParamVO param) {
        //查询需要重新训练的图片
        List<SceneMarkShape> reDetectStatuList = findByReTrainStatus(1);
        for (SceneMarkShape shape : reDetectStatuList) {
            SceneMarkShapeReDetectParamVO paramVO=new SceneMarkShapeReDetectParamVO();
            paramVO.setWebSite(mainUrl);
            if (StrUtil.isNotEmpty(param.getSaveDir())){
                paramVO.setSaveDir(param.getSaveDir());
            }
            paramVO.setNum(shape.getNum());
            paramVO.setImagePath(shape.getImagePath());
            paramVO.setDetectType(1);
            rabbitMqProducer.sendByWorkQueue(yolov5TrainQueue,paramVO);
        }
        //查询需要进入训练的场景
        List<SceneMarkShape> byToDetectStatus = findByToTrainStatus(0);
        for (SceneMarkShape shape : byToDetectStatus) {
            SceneMarkShapeReDetectParamVO paramVO=new SceneMarkShapeReDetectParamVO();
            paramVO.setWebSite(mainUrl);
            if (StrUtil.isNotEmpty(param.getSaveDir())){
                paramVO.setSaveDir(param.getSaveDir());
            }
            paramVO.setNum(shape.getNum());
            paramVO.setDetectType(2);
            rabbitMqProducer.sendByWorkQueue(yolov5TrainQueue,paramVO);
        }
    }

    @Override
    public ResultData editLabelByFile(String num, String imgPath, MultipartFile file) throws IOException {

        SceneMarkShape shape = findByNumAndImagePath(num, imgPath);
        if (ObjectUtil.isNotNull(shape)){
            String uuid = UUID.randomUUID().toString();
            String fileName = file.getOriginalFilename();
            String extName = FileUtil.extName(fileName);
            String tempFileName = uuid + "." + extName;
            String srcPath = ConstantFilePath.SCENE_V4_PATH + num + "/markShapes/" + tempFileName;
            File tempFile = new File(srcPath);
            if(!tempFile.getParentFile().exists()){
                tempFile.getParentFile().mkdirs();
            }
            file.transferTo(tempFile);

            List<String> s = FileUtil.readUtf8Lines(tempFile);
            List<JSONObject> shapeJsons=new ArrayList<>();
            //转换labelimg标注处理的结果
            log.info("转换labelimg标注处理的结果开始");
            for (String s1 : s) {
                int[] ints = ConverxyUtil.centerWh2xyxy(s1, 4096,2048);
                String[] s2 = s1.split(" ");
                JSONObject shapeJson=new JSONObject();
                shapeJson.put("bbox",ints);
                shapeJson.put("color",ConverxyUtil.getColor(s2[0]));
                shapeJson.put("label",s1);
                shapeJson.put("category", ConverxyUtil.getLabelVal(s2[0]));
                shapeJson.put("score",0);
                shapeJsons.add(shapeJson);
            }
            log.info("转换labelimg标注处理的结果结束，{}",shapeJsons);
            shape.setShapes(shapeJsons);
            shape.setReTrain(1);
            shape.setUpdateTime(new Date());
            updateById(shape);
            return ResultData.ok(shape);
        }else {
            return ResultData.error(ErrorCode.NOT_RECORD);
        }
    }

    @Override
    public SceneMarkShape findByNumAndImagePath(String num, String imagePath) {
        LambdaQueryWrapper<SceneMarkShape> wrapper = Wrappers.lambdaQuery();
        wrapper.eq(SceneMarkShape::getNum,num);
        wrapper.eq(SceneMarkShape::getImagePath,imagePath);
        return getOne(wrapper);
    }

    @Override
    public void saveFileToDB(MultipartFile file, String num) throws IOException {
        String uuid = UUID.randomUUID().toString();
        String fileName = file.getOriginalFilename();
        String extName = FileUtil.extName(fileName);
        String tempFileName = uuid + "." + extName;
        String srcPath = ConstantFilePath.SCENE_V4_PATH + num + "/markShapes/" + tempFileName;
        File tempFile = new File(srcPath);
        if(!tempFile.getParentFile().exists()){
            tempFile.getParentFile().mkdirs();
        }
        file.transferTo(tempFile);
        String s = FileUtil.readUtf8String(tempFile);
        JSONObject jsonObject = JSONObject.parseObject(s);
        tempFile.delete();
        SceneMarkShape sceneMarkShape= JSON.toJavaObject(jsonObject,SceneMarkShape.class);
        sceneMarkShape.setNum(num);
        SceneMarkShape shape = findByNumAndImagePath(sceneMarkShape.getNum(), sceneMarkShape.getImagePath());
        if (ObjectUtil.isNotNull(shape)){
            log.info("shape-替换id修改---{}",sceneMarkShape);
            sceneMarkShape.setId(shape.getId());
            updateById(sceneMarkShape);
        }else {
            log.info("新增-替换id修改---{}",sceneMarkShape);
            log.info("MarkShapeMapper---{}",sceneMarkShape);
            save(sceneMarkShape);
        }
    }

    @Override
    public List<SceneMarkShape> findByNum(String num) {
        LambdaQueryWrapper<SceneMarkShape> wrapper = Wrappers.lambdaQuery();
        wrapper.eq(SceneMarkShape::getNum,num);
        return list(wrapper);
    }

}