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

package com.fdkankan.scene.service.impl;

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.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.model.constants.ConstantFilePath;
import com.fdkankan.rabbitmq.util.RabbitMqProducer;
import com.fdkankan.scene.entity.SceneMarkShape;
import com.fdkankan.scene.mapper.MarkShapeMapper;
import com.fdkankan.scene.service.ISceneMarkShapeService;
import com.fdkankan.scene.util.ConverxyUtil;
import com.fdkankan.scene.vo.SceneMarkShapeParamVO;
import com.fdkankan.scene.vo.SceneMarkShapeReDetectParamVO;
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.ArrayList;
import java.util.Date;
import java.util.List;


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

    @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 = cn.hutool.core.io.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> findByNumAndType(String num,Integer type) {
        LambdaQueryWrapper<SceneMarkShape> wrapper = Wrappers.lambdaQuery();
        wrapper.eq(SceneMarkShape::getNum,num);
        wrapper.eq(SceneMarkShape::getType,type);
        return list(wrapper);
    }

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