Jelajahi Sumber

feat: 目标检测优化搞定

WANGKANG 8 bulan lalu
induk
melakukan
f87140d3d4

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

@@ -9,6 +9,7 @@ public class BizConstant {
     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 TD_TRAIN_MODEL_SUFFIX = ".ckpt";
 
     public static final String UNZIP_SUFFIX = "_unzip";
     public static final String STABLE_SUFFIX = "_stable";

+ 5 - 0
taais-modules/taais-biz/src/main/java/com/taais/biz/controller/TargetDetectionController.java

@@ -130,4 +130,9 @@ public class TargetDetectionController extends BaseController {
     public CommonResult getLog(@PathVariable("id") Long id) {
         return targetDetectionService.getLog(id);
     }
+
+    @GetMapping("/model/{id}")
+    public CommonResult getModelList(@PathVariable("id") Long id) {
+        return targetDetectionService.getModelList(id);
+    }
 }

+ 2 - 0
taais-modules/taais-biz/src/main/java/com/taais/biz/service/ITargetDetectionService.java

@@ -76,4 +76,6 @@ public interface ITargetDetectionService extends IBaseService<TargetDetection> {
     ResponseEntity<Resource> zipImages(Long id);
 
     CommonResult getLog(Long id);
+
+    CommonResult getModelList(Long id);
 }

+ 2 - 2
taais-modules/taais-biz/src/main/java/com/taais/biz/service/impl/AlgorithmModelTrackServiceImpl.java

@@ -183,8 +183,8 @@ public class AlgorithmModelTrackServiceImpl extends BaseServiceImpl<AlgorithmMod
         } // 这是使用ossid上传
         else if(ObjectUtil.isNotEmpty(algorithmModelTrackBo.getModelPath())) { // 这是使用路径上传
             File modelFile = new File(algorithmModelTrackBo.getModelPath());
-            if (!modelFile.exists() || !modelFile.isFile()) {
-                return CommonResult.fail("请指定正确的模型路径!");
+            if (!modelFile.exists()) {
+                return CommonResult.fail("模型路径不存在!");
             }
             algorithmModelTrack.setModelAddress(algorithmModelTrackBo.getModelPath());
         }

+ 57 - 4
taais-modules/taais-biz/src/main/java/com/taais/biz/service/impl/TargetDetectionServiceImpl.java

@@ -7,10 +7,8 @@ import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.util.Arrays;
-import java.util.Date;
-import java.util.HashMap;
-import java.util.List;
+import java.text.DecimalFormat;
+import java.util.*;
 
 import cn.hutool.core.util.ObjectUtil;
 import com.mybatisflex.core.paginate.Page;
@@ -61,6 +59,9 @@ import com.alibaba.fastjson2.JSON;
  */
 @Service
 public class TargetDetectionServiceImpl extends BaseServiceImpl<TargetDetectionMapper, TargetDetection> implements ITargetDetectionService {
+    @Value("${server.port}")
+    String port;
+
     @Value("${server.task_stop_url}")
     private String task_stop_url;
 
@@ -427,4 +428,56 @@ public class TargetDetectionServiceImpl extends BaseServiceImpl<TargetDetectionM
         }
         return CommonResult.success(readLogContent(logPath), "success");
     }
+
+    @Override
+    public CommonResult getModelList(Long id) {
+        TargetDetection entity = getById(id);
+        String outPutPath = entity.getOutputPath();
+        String modelPath = outPutPath;
+        File modelDir = new File(modelPath);
+        if (!modelDir.exists()) {
+            return CommonResult.fail("模型输出目录不存在!");
+        }
+        File[] files = modelDir.listFiles();
+        Integer idx = 0;
+        ArrayList<Map<String, String>> res = new ArrayList<>();
+        for (File file : files) {
+            if(!file.getName().endsWith(BizConstant.TD_TRAIN_MODEL_SUFFIX)){
+                continue;
+            }
+            idx += 1;
+            Map<String, String> tmp = new HashMap<>();
+            tmp.put("id", idx.toString());
+            tmp.put("name", file.getName());
+            String filePath = file.getPath();
+            tmp.put("path", filePath);
+            // todo 获取真正的url
+            // http://localhost:9090/profile/upload/2024/10/27/1_1729404909511_20241027153840A001.zip
+            String url = "http://localhost:" + port + Constants.RESOURCE_PREFIX + filePath.substring(TaaisConfig.getProfile().length());
+            url = url.replaceAll("\\\\", "/"); // windows
+            tmp.put("url", url);
+
+            double fileSize = (getFileSize(file) / (1024.0 * 1024.0));
+            DecimalFormat decimalFormat = new DecimalFormat("#.##");
+            String formatFileSize = decimalFormat.format(fileSize);
+            tmp.put("size", formatFileSize + "MB");
+            res.add(tmp);
+        }
+        return CommonResult.success(res, "success");
+    }
+    public static long getFileSize(File file) {
+        // 返回单位B
+        if(file.isFile()) {
+            return file.length();
+        }
+        else {
+            long totalSize = 0;
+            if(file.listFiles() != null && file.listFiles().length > 0) {
+                for(File file__ : Objects.requireNonNull(file.listFiles())) {
+                    totalSize += getFileSize(file__);
+                }
+            }
+            return totalSize;
+        }
+    }
 }

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

@@ -417,8 +417,8 @@ public class ToInfraredServiceImpl extends BaseServiceImpl<ToInfraredMapper, ToI
     public CommonResult getModelList(Long id) {
         ToInfrared entity = getById(id);
         String outPutPath = entity.getOutputPath();
-        String logPath = outPutPath + File.separator + "model";
-        File modelDir = new File(logPath);
+        String modelPath = outPutPath + File.separator + "model";
+        File modelDir = new File(modelPath);
         if (!modelDir.exists()) {
             return CommonResult.fail("模型输出目录不存在!");
         }