diff --git a/src/main/java/com/kamco/cd/kamcoback/model/service/ModelMngService.java b/src/main/java/com/kamco/cd/kamcoback/model/service/ModelMngService.java index f7447f00..14d2a9af 100644 --- a/src/main/java/com/kamco/cd/kamcoback/model/service/ModelMngService.java +++ b/src/main/java/com/kamco/cd/kamcoback/model/service/ModelMngService.java @@ -15,6 +15,7 @@ import com.kamco.cd.kamcoback.upload.service.UploadService; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; +import java.nio.file.Paths; import java.time.LocalDate; import java.util.List; import java.util.UUID; @@ -237,22 +238,43 @@ public class ModelMngService { FIleChecker.getFilesFromAllDepth( dirPath, "*", "pth,py,pt,json", 10, "name", startPos, endPos); + boolean hasPt = false; // pt 파일 존재 여부 + for (Basic dto : files) { // 예: 파일명 출력 및 추가 작업 String foldNm = dto.getFullPath().replace(dto.getFileNm(), ""); - if (dto.getExtension().equals("pth")) { - modelUploadResDto.setCdModelPath(foldNm); - modelUploadResDto.setCdModelFileName(dto.getFileNm()); - } else if (dto.getExtension().equals("py")) { - modelUploadResDto.setCdModelConfigPath(foldNm); - modelUploadResDto.setCdModelConfigFileName(dto.getFileNm()); - } else if (dto.getExtension().equals("pt")) { - modelUploadResDto.setClsModelPath(foldNm); - modelUploadResDto.setClsModelFileName(dto.getFileNm()); - } else if (dto.getExtension().equals("json")) { - modelUploadResDto.setJsonPath(foldNm); - modelUploadResDto.setJsonFileName(dto.getFileNm()); + switch (dto.getExtension()) { + case "pth" -> { + modelUploadResDto.setCdModelPath(foldNm); + modelUploadResDto.setCdModelFileName(dto.getFileNm()); + } + case "py" -> { + modelUploadResDto.setCdModelConfigPath(foldNm); + modelUploadResDto.setCdModelConfigFileName(dto.getFileNm()); + } + case "pt" -> { + modelUploadResDto.setClsModelPath(foldNm); + modelUploadResDto.setClsModelFileName(dto.getFileNm()); + hasPt = true; + } + case "json" -> { + modelUploadResDto.setJsonPath(foldNm); + modelUploadResDto.setJsonFileName(dto.getFileNm()); + } + } + } + + if (!hasPt) { + + String defaultPath = "/kamco-nfs/ckpt/classification/"; + String defaultFileName = "v5-best.pt"; + + Path ptPath = Paths.get(defaultPath, defaultFileName); + + if (Files.exists(ptPath)) { + modelUploadResDto.setClsModelPath(defaultPath); + modelUploadResDto.setClsModelFileName(defaultFileName); } }