Browse Source

feat: 可见光转红外对接完成

WANGKANG 10 months ago
parent
commit
a135b4249f

+ 2 - 1
taais-modules/taais-biz/src/main/java/com/taais/biz/constant/BizConstant.java

@@ -6,8 +6,9 @@ package com.taais.biz.constant;
 public class BizConstant {
     // 视频文件存储目录
     public static final String UPLOAD_DIR = "uploads/";
-    public static final String MODEL_PATH = "model/";
+    public static final String MODEL_PATH = "model";
     public static final String MODEL_SUFFIX = ".pt";
+    public static final String PREDICT_PATH = "predict";
 
     public static final String UNZIP_SUFFIX = "_unzip";
     public static final String TO_INFRARED_SUFFIX = "_to_infrared";

+ 48 - 9
taais-modules/taais-biz/src/main/java/com/taais/biz/service/impl/ToInfraredServiceImpl.java

@@ -268,14 +268,14 @@ public class ToInfraredServiceImpl extends BaseServiceImpl<ToInfraredMapper, ToI
 //            startToInfraredTask.setSource_dir("C:" + toInfrared.getInputPath());
 //            startToInfraredTask.setResult_dir("C:" + toInfrared.getOutputPath());
 //        } else {
-            startToInfraredTask.setSource_dir(toInfrared.getInputPath());
-            startToInfraredTask.setResult_dir(toInfrared.getOutputPath());
+        startToInfraredTask.setSource_dir(toInfrared.getInputPath());
+        startToInfraredTask.setResult_dir(toInfrared.getOutputPath());
 //        }
 
         // startToInfraredTask.setSource_dir(toInfrared.getInputPath());
         // startToInfraredTask.setResult_dir(toInfrared.getOutputPath());
 
-        if ( BizConstant.AlgorithmType.REASONING.equals(algorithmConfigTrack.getType())) {
+        if (BizConstant.AlgorithmType.REASONING.equals(algorithmConfigTrack.getType())) {
             startToInfraredTask.setModel_path(algorithmModelTrack.getModelAddress());
         }
 
@@ -359,29 +359,68 @@ public class ToInfraredServiceImpl extends BaseServiceImpl<ToInfraredMapper, ToI
             return ResponseEntity.status(HttpStatus.NOT_FOUND).body(null);
         }
 
-        String outputPath = toInfrared.getOutputPath();
-
 //        String osName = System.getProperty("os.name");
 //        if (osName.toLowerCase().contains("windows")) {
 //            outputPath = outputPath;
 //        }
 
-        File  modelDir= new File(outputPath + File.separator + BizConstant.MODEL_PATH);
+        AlgorithmModelTrackVo model = algorithmModelTrackService.selectById(toInfrared.getAlgorithmModelId());
+        AlgorithmConfigTrackVo config = algorithmConfigTrackService.selectById(model.getAlgorithmId());
+
+        if (BizConstant.AlgorithmType.TRAIN.equals(config.getType())) { // 如果是训练类型,返回模型
+            return getLastModelFile(toInfrared);
+        } else if (BizConstant.AlgorithmType.REASONING.equals(config.getType())) { // 如果是推理类型,返回图片集
+            return getPredictImages(toInfrared);
+        } else {
+            return ResponseEntity.status(HttpStatus.NOT_FOUND).body(null);
+        }
+    }
+
+    private ResponseEntity<org.springframework.core.io.Resource> getPredictImages(ToInfrared toInfrared) {
+        String outputPath = toInfrared.getOutputPath();
+        String predictDir = outputPath + File.separator + BizConstant.PREDICT_PATH;
+
+        String zipFilePath = outputPath + File.separator + "predict.zip";
+
+        File file = new File(zipFilePath);
+        if (!file.exists()) {
+            try {
+                ZipUtils.zipFolderFiles(predictDir, zipFilePath);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+
+        if (!file.exists() || !file.isFile()) {
+            return ResponseEntity.status(HttpStatus.NOT_FOUND).body(null);
+        }
+
+        org.springframework.core.io.Resource resource = new FileSystemResource(file);
+        return ResponseEntity.ok()
+            .header(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"" + file.getName() + "\"")
+            .header(HttpHeaders.CONTENT_TYPE, "application/octet-stream")
+            .body(resource);
+    }
+
+    public ResponseEntity<org.springframework.core.io.Resource> getLastModelFile(ToInfrared toInfrared) {
+        String outputPath = toInfrared.getOutputPath();
+        File modelDir = new File(outputPath + File.separator + BizConstant.MODEL_PATH);
 
         if (!modelDir.exists() || !modelDir.isDirectory()) {
             return ResponseEntity.status(HttpStatus.NOT_FOUND).body(null);
         }
 
         File[] modelFiles = modelDir.listFiles();
-        if(modelFiles == null || modelFiles.length == 0) {
+        if (modelFiles == null || modelFiles.length == 0) {
             return ResponseEntity.status(HttpStatus.NOT_FOUND).body(null);
         }
 
-        for(int i=0; i< modelFiles.length; i++) {
+        for (int i = 0; i < modelFiles.length; i++) {
             System.out.println(modelFiles[i].getName());
         }
 
-        File returnFile = modelFiles[modelFiles.length-1];
+        File returnFile = modelFiles[modelFiles.length - 1];
 
         org.springframework.core.io.Resource resource = new FileSystemResource(returnFile);
         return ResponseEntity.ok()