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 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 findByReTrainStatus(Integer reTrain){ LambdaQueryWrapper wrapper = Wrappers.lambdaQuery(); wrapper.eq(SceneMarkShape::getReTrain,reTrain); return list(wrapper); } @Override public List 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 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 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 s = FileUtil.readUtf8Lines(tempFile); List 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 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 findByNumAndType(String num,Integer type) { LambdaQueryWrapper wrapper = Wrappers.lambdaQuery(); wrapper.eq(SceneMarkShape::getNum,num); wrapper.eq(SceneMarkShape::getType,type); return list(wrapper); } @Override public List findByNum(String num) { LambdaQueryWrapper wrapper = Wrappers.lambdaQuery(); wrapper.eq(SceneMarkShape::getNum,num); return list(wrapper); } }