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 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 imagePaths = param.getList().stream().map(v -> v.getImagePath()).collect(Collectors.toSet()); List sceneMarkShapes = this.baseMapper.selectByNumAndImagePaths(param.getNum(), imagePaths); Map imagePathIdMap = sceneMarkShapes.stream().collect(Collectors.toMap(SceneMarkShape::getImagePath, SceneMarkShape::getId)); List 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 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 = 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 findByNum(String num) { LambdaQueryWrapper wrapper = Wrappers.lambdaQuery(); wrapper.eq(SceneMarkShape::getNum,num); return list(wrapper); } }