feat: result export

taais-modules/taais-biz/src/main/java/com/taais/biz/controller/DataAugmentationController.java

@@ -1,6 +1,8 @@
 package com.taais.biz.controller;
 
 
+import java.io.File;
+import java.io.FileInputStream;
 import java.io.IOException;
 import java.nio.file.DirectoryStream;
 import java.nio.file.Files;
@@ -13,6 +15,7 @@ import java.util.stream.Stream;
 import com.mybatisflex.core.query.QueryWrapper;
 import com.taais.biz.domain.DataAugmentation;
 import com.taais.biz.domain.bo.DataAugmentationBo;
+import com.taais.biz.domain.bo.DataAugmentationResultBo;
 import com.taais.biz.domain.bo.VideoStableStartResultBo;
 import com.taais.biz.domain.vo.DataAugmentationVo;
 import com.taais.biz.service.IVideoStableService;
@@ -23,7 +26,10 @@ import jakarta.validation.Valid;
 import lombok.RequiredArgsConstructor;
 import jakarta.servlet.http.HttpServletResponse;
 import cn.dev33.satoken.annotation.SaCheckPermission;
+import org.apache.commons.io.IOUtils;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.MediaType;
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.*;
 import org.springframework.validation.annotation.Validated;
@@ -96,7 +102,7 @@ public class DataAugmentationController extends BaseController {
             String fileName = imagePath.getFileName().toString();
             Map<String, List<String>> images = new HashMap<>();
             //图像拼接算法有多个输入图片
-            if ("图像拼接_sift".equals(dataAugmentation.getTaskType()) || "图像拼接_coordinate".equals(dataAugmentation.getTaskType())) {
+            if ("侦察图像拼接算法_sift".equals(dataAugmentation.getTaskType()) || "侦察图像拼接算法_coordinate".equals(dataAugmentation.getTaskType())) {
                 String lastDirectoryName = imagePath.getFileName().toString();
                 outputPath = outputPath.resolve(lastDirectoryName);  //得到推理结果目录
                 Stream<Path> stream = Files.list(outputPath);
@@ -156,8 +162,8 @@ public class DataAugmentationController extends BaseController {
     }
 
     @PostMapping("/get_result")
-    public CommonResult getResult(@Valid @RequestBody VideoStableStartResultBo videoStableStartResultBo) {
-        return dataAugmentationService.getResult(videoStableStartResultBo);
+    public CommonResult getResult(@Valid @RequestBody DataAugmentationResultBo dataAugmentationResultBo) {
+        return dataAugmentationService.getResult(dataAugmentationResultBo);
     }
 
     @GetMapping("/start/{id}")
@@ -185,9 +191,30 @@ public class DataAugmentationController extends BaseController {
     @SaCheckPermission("demo:dataAugmentation:export")
     @Log(title = "数据增强", businessType = BusinessType.EXPORT)
     @PostMapping("/export")
-    public void export(HttpServletResponse response, DataAugmentationBo dataAugmentationBo) {
-        List<DataAugmentationVo> list = dataAugmentationService.selectList(dataAugmentationBo);
-        ExcelUtil.exportExcel(list, "数据增强", DataAugmentationVo.class, response);
+    public void export(HttpServletResponse response, @RequestBody String zipDirPath) {
+
+        try {
+            // 假设你已经知道ZIP文件的路径
+//            String zipFilePath = dataAugmentationVo.getOutputPath() + ".zip";
+//            System.out.println(dataAugmentationVo.toString());
+
+            String zipFilePath = zipDirPath.substring(1, zipDirPath.length() - 1) + ".zip";
+            System.out.println( zipFilePath);
+            File zipFile = new File(zipFilePath);
+            String fileName = zipFile.getName();
+            // 设置响应头
+            response.setContentType(MediaType.APPLICATION_OCTET_STREAM_VALUE);
+            response.setHeader(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=" + fileName);
+            // 读取ZIP文件并写入响应流
+            try (FileInputStream fis = new FileInputStream(zipFilePath)) {
+                IOUtils.copy(fis, response.getOutputStream());
+            }
+            // 刷新输出流
+            response.flushBuffer();
+        } catch (IOException e) {
+            e.printStackTrace();
+            response.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
+        }
     }
 
     /**

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

@@ -3,12 +3,7 @@ package com.taais.biz.controller;
 import com.alibaba.fastjson2.JSON;
 import com.taais.biz.constant.BizConstant;
 import com.taais.biz.domain.TaskTrackResultBo;
-import com.taais.biz.domain.bo.AlgorithmModelBo;
-import com.taais.biz.domain.bo.ObjectTraceMergeBo;
-import com.taais.biz.domain.bo.TargetIdentificationSubtaskDetailsBo;
-import com.taais.biz.domain.bo.ToInfraredBo;
-import com.taais.biz.domain.bo.ToInfraredStartResultBo;
-import com.taais.biz.domain.bo.VideoStableStartResultBo;
+import com.taais.biz.domain.bo.*;
 import com.taais.biz.domain.dto.TaskResultDTO;
 import com.taais.biz.service.*;
 import com.taais.biz.service.impl.AlgorithmModelServiceImpl;
@@ -150,7 +145,7 @@ public class PublicController extends BaseController {
         return videoStableService.getResult(videoStableStartResultBo);
     }
     @PostMapping("/dataAugmentation/get_result")
-    public CommonResult Result(@Valid @RequestBody VideoStableStartResultBo dataAugmentationStartResultBo) {
+    public CommonResult Result(@Valid @RequestBody DataAugmentationResultBo dataAugmentationStartResultBo) {
         return dataAugmentationService.getResult(dataAugmentationStartResultBo);
     }
 

taais-modules/taais-biz/src/main/java/com/taais/biz/domain/bo/DataAugmentationResultBo.java

@@ -0,0 +1,25 @@
+package com.taais.biz.domain.bo;
+
+import jakarta.validation.constraints.NotNull;
+import lombok.*;
+
+/**
+ * @Datetime : 2023/12/5 12:13
+ * @Author : WANGKANG
+ * @Email : 1686617586@qq.com
+ * @File : VideoStableStartResult.java
+ * @Brief :
+ */
+
+@Builder
+@Data
+@NoArgsConstructor
+@AllArgsConstructor
+@EqualsAndHashCode()
+public class DataAugmentationResultBo {
+    @NotNull(message = "status不能为空")
+    private String status;
+    private String msg;
+    @NotNull(message = "id不能为空")
+    private Long id;
+}

taais-modules/taais-biz/src/main/java/com/taais/biz/service/IDataAugmentationService.java

@@ -2,6 +2,7 @@ package com.taais.biz.service;
 
 import com.taais.biz.domain.DataAugmentation;
 import com.taais.biz.domain.bo.DataAugmentationBo;
+import com.taais.biz.domain.bo.DataAugmentationResultBo;
 import com.taais.biz.domain.bo.DataAugmentationStartResultBo;
 import com.taais.biz.domain.bo.VideoStableStartResultBo;
 import com.taais.biz.domain.vo.DataAugmentationVo;
@@ -72,7 +73,7 @@ public interface IDataAugmentationService extends IBaseService<DataAugmentation>
 
     CommonResult start(Long id);
 
-    CommonResult getResult(VideoStableStartResultBo videoStableStartResultBo);
+    CommonResult getResult(DataAugmentationResultBo videoStableStartResultBo);
 
     Object getCompareNum(Long taskId);
 

taais-modules/taais-biz/src/main/java/com/taais/biz/service/impl/DataAugmentationServiceImpl.java

@@ -8,6 +8,7 @@ import com.taais.biz.domain.HttpResponseEntity;
 import com.taais.biz.domain.DataAugmentation;
 import com.taais.biz.domain.TransmissionObject;
 import com.taais.biz.domain.bo.DataAugmentationBo;
+import com.taais.biz.domain.bo.DataAugmentationResultBo;
 import com.taais.biz.domain.bo.DataAugmentationStartResultBo;
 import com.taais.biz.domain.bo.VideoStableStartResultBo;
 import com.taais.biz.domain.vo.DataAugmentationVo;
@@ -46,15 +47,15 @@ import org.springframework.transaction.annotation.Transactional;
 
 import java.io.File;
 import java.io.IOException;
-import java.nio.file.DirectoryStream;
-import java.nio.file.Files;
-import java.nio.file.Path;
-import java.nio.file.Paths;
+import java.nio.file.*;
+import java.nio.file.attribute.BasicFileAttributes;
 import java.util.Arrays;
 import java.util.Date;
 import java.util.HashMap;
 import java.util.List;
 import java.util.stream.Stream;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipOutputStream;
 
 import static com.taais.biz.domain.table.DataAugmentationTableDef.DATA_AUGMENTATION;
 /**
@@ -314,22 +315,51 @@ public class DataAugmentationServiceImpl extends BaseServiceImpl<DataAugmentatio
         }
     }
 
+    private static void createZipArchive(Path sourceDir, Path zipFilePath) throws IOException {
+        try (ZipOutputStream zos = new ZipOutputStream(Files.newOutputStream(zipFilePath))) {
+            Files.walkFileTree(sourceDir, new SimpleFileVisitor<Path>() {
+                public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
+                    Path targetFile = sourceDir.relativize(file);
+                    ZipEntry zipEntry = new ZipEntry(targetFile.toString());
+                    zipEntry.setTime(attrs.lastModifiedTime().toMillis());
+                    zos.putNextEntry(zipEntry);
+                    byte[] bytes = Files.readAllBytes(file);
+                    zos.write(bytes, 0, bytes.length);
+                    zos.closeEntry();
+                    return FileVisitResult.CONTINUE;
+                }
+
+                public FileVisitResult preVisitDirectory(Path dir, BasicFileAttributes attrs) throws IOException {
+                    // 不需要为目录添加条目，除非你想记录空目录
+                    return FileVisitResult.CONTINUE;
+                }
+            });
+        }
+    }
     @Override
-    public CommonResult getResult(VideoStableStartResultBo dataAugmentationStartResultBo) {
+    public CommonResult getResult(DataAugmentationResultBo dataAugmentationStartResultBo) {
         Long id = dataAugmentationStartResultBo.getId();
         String status = dataAugmentationStartResultBo.getStatus();
         String msg = dataAugmentationStartResultBo.getMsg();
         DataAugmentation dataAugmentation = getById(id);
+        String outputPath = dataAugmentation.getOutputPath();
+        // 要压缩的目录路径
+        Path sourceDir = Paths.get(outputPath);
+        // 输出的ZIP文件路径
+        Path zipFilePath = sourceDir.resolveSibling(sourceDir.getFileName() + ".zip");
         dataAugmentation.setLog(msg);
         dataAugmentation.setStatus("200".equals(status) ? BizConstant.ModelStatus.END : BizConstant.ModelStatus.FAILED);
         dataAugmentation.setEndTime(new Date());
+        dataAugmentation.setCostSecond((dataAugmentation.getEndTime().getTime() - dataAugmentation.getStartTime().getTime()) / 1000);
         try {
-            dataAugmentation.setCostSecond((dataAugmentation.getEndTime().getTime() - dataAugmentation.getStartTime().getTime()) / 1000);
+            updateById(dataAugmentation);
+            createZipArchive(sourceDir, zipFilePath);
+            System.out.println("Directory has been zipped successfully: " + zipFilePath);
+            return CommonResult.success();
         } catch (Exception e) {
-            dataAugmentation.setCostSecond(null);
+            System.out.println("Directory has not been zipped successfully: " + zipFilePath);
+            return  CommonResult.fail("Directory has not been zipped successfully: " + zipFilePath);
         }
-        updateById(dataAugmentation);
-        return CommonResult.success();
     }
 
     @Override