Forráskód Böngészése

[FIX] fix model save issue

Eagle 7 hónapja
szülő
commit
21ed1177a7

+ 15 - 1
taais-modules/taais-biz/src/main/java/com/taais/biz/controller/PublicController.java

@@ -10,6 +10,7 @@ import com.taais.biz.service.impl.AlgorithmModelServiceImpl;
 import com.taais.biz.service.impl.ObjectTraceMergeServiceImpl;
 import com.taais.biz.service.impl.TargetIdentificationSubtaskDetailsServiceImpl;
 import com.taais.biz.service.service.impl.ObjectMatchServiceImpl;
+import com.taais.common.core.config.TaaisConfig;
 import com.taais.common.core.core.domain.CommonResult;
 import com.taais.common.log.annotation.Log;
 import com.taais.common.log.enums.BusinessType;
@@ -24,6 +25,7 @@ import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.validation.annotation.Validated;
 import org.springframework.web.bind.annotation.*;
 
+import java.io.File;
 import java.util.Arrays;
 import java.util.Date;
 
@@ -96,7 +98,19 @@ public class PublicController extends BaseController {
                 Long algorithmId = detailsBo.getAlgorithmId();
                 AlgorithmModelBo bo = new AlgorithmModelBo();
                 bo.setAlgorithmId(algorithmId);
-                bo.setModelAddress("/profile/task" + detailsBo.getResultPath() + "weights/best.pt");
+                String _path = TaaisConfig.getProfile() + "/task" + detailsBo.getResultPath() + "/weights";
+                File dir = new File(_path);
+                String SUFFIX_NAME = "NONAME";
+                if (dir.exists()) {
+                    File[] files = dir.listFiles();
+                    for (File file : files) {
+                        if (file.isFile() && file.getName().contains("best")) {
+                            SUFFIX_NAME = file.getName();
+                            break;
+                        }
+                    }
+                }
+                bo.setModelAddress("/profile/task" + detailsBo.getResultPath() + "/weights/" + SUFFIX_NAME);
                 bo.setModelName(detailsBo.getName() + "_" + detailsBo.getCreateTime().toString());
                 algorithmModelService.insert(bo);
             }