server/distributed-print/src/test/java/com/qmth/distributed/print/ExportAiResult.java

package com.qmth.distributed.print;

import com.alibaba.fastjson.JSON;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.qmth.boot.api.exception.ApiException;
import com.qmth.boot.tools.excel.ExcelWriter;
import com.qmth.boot.tools.excel.enums.ExcelType;
import com.qmth.distributed.print.dto.ExportAiResultDto;
import com.qmth.teachcloud.common.bean.dto.mark.PictureConfig;
import com.qmth.teachcloud.common.bean.vo.FilePathVo;
import com.qmth.teachcloud.common.contant.SystemConstant;
import com.qmth.teachcloud.common.entity.MarkQuestion;
import com.qmth.teachcloud.common.enums.ExceptionResultEnum;
import com.qmth.teachcloud.common.service.FileUploadService;
import com.qmth.teachcloud.common.util.FileUtil;
import com.qmth.teachcloud.mark.dto.ai.SheetImageDto;
import com.qmth.teachcloud.mark.entity.MarkOcrStudentQuestion;
import com.qmth.teachcloud.mark.entity.MarkStudent;
import com.qmth.teachcloud.mark.entity.MarkTask;
import com.qmth.teachcloud.mark.enums.MarkTaskStatus;
import com.qmth.teachcloud.mark.service.MarkOcrStudentQuestionService;
import com.qmth.teachcloud.mark.service.MarkQuestionService;
import com.qmth.teachcloud.mark.service.MarkStudentService;
import com.qmth.teachcloud.mark.service.MarkTaskService;
import com.qmth.teachcloud.mark.utils.AiUtil;
import org.apache.commons.collections4.CollectionUtils;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;

import javax.annotation.Resource;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 根据名单，导出测试环境AI阅卷结果
 */
@SpringBootTest
@RunWith(SpringRunner.class)
public class ExportAiResult {
    @Resource
    private MarkOcrStudentQuestionService markOcrStudentQuestionService;
    @Resource
    private MarkStudentService markStudentService;
    @Resource
    private MarkQuestionService markQuestionService;
    @Resource
    private MarkTaskService markTaskService;
    @Resource
    private FileUploadService fileUploadService;

    private String rootPath = "/Users/xiaofei/aaaa";

    @Test
    public void export() {
        long questionId = 657572149854208079l;
        MarkQuestion markQuestion = markQuestionService.getById(questionId);

        List<MarkTask> markTaskList = listMarkTaskByQuestionId(questionId);

        List<ExportAiResultDto> list = new ArrayList<>();
        for (MarkTask markTask : markTaskList) {
            ExportAiResultDto dto = new ExportAiResultDto();
            // 识别区域剪切
            getOcrRegion(markTask, markQuestion, dto);
            list.add(dto);
        }

        export(list);
    }

    private void getOcrRegion(MarkTask markTask, MarkQuestion markQuestion, ExportAiResultDto dto) {
        List<File> deleteFileList = new ArrayList<>();
        Long studentId = markTask.getStudentId();
        // 读取考生图片
        MarkStudent markStudent = markStudentService.getById(studentId);
        dto.setStudentCode(markStudent.getStudentCode());
        dto.setMarkerScore(markTask.getMarkerScore());

        Map<Integer, SheetImageDto> fileMap = new HashMap<>();
        try {
            //考生题卡原图
            List<FilePathVo> filePathVoList = markStudent.listSheetPath();
            for (int i = 0; i < filePathVoList.size(); i++) {
                FilePathVo filePathVo = filePathVoList.get(i);
                File file = new File(rootPath + File.separator + "sheet" + File.separator + markStudent.getStudentCode() + "-" + (i + 1) + SystemConstant.JPG_PREFIX);
                try {
                    if (!file.exists()) {
                        file.getParentFile().mkdirs();
                        file.createNewFile();
                    }
                    fileUploadService.downloadFile(filePathVo.getType(), filePathVo.getUploadType(), filePathVo.getPath(), file.getPath());
                } catch (Exception e) {
                    return;
//                    throw ExceptionResultEnum.ERROR.exception("读取考生原图失败" + "，考号[" + markStudent.getStudentCode() + "。" + e.getMessage());
                }
                fileMap.put(i + 1, new SheetImageDto(i + 1, file, filePathVo.getMd5(), AiUtil.imageDim(file)));
                deleteFileList.add(file);
            }

            // 题目评卷区
            List<PictureConfig> pictureConfigList = JSON.parseArray(markQuestion.getPicList(), PictureConfig.class);

            for (int i = 0; i < pictureConfigList.size(); i++) {
                PictureConfig pictureConfig = pictureConfigList.get(i);
                SheetImageDto sheetImageDto = fileMap.get(pictureConfig.getI());
                File file = new File(rootPath + File.separator + (MarkTaskStatus.MARKED.equals(markTask.getStatus()) ? "cut" : "error") + File.separator + markStudent.getStudentCode() + "-" + (i + 1) + SystemConstant.JPG_PREFIX);
                if (!file.exists()) {
                    file.getParentFile().mkdirs();
                    file.createNewFile();
                }
                AiUtil.subImg(sheetImageDto, pictureConfig, file);
            }

            List<MarkOcrStudentQuestion> markOcrStudentQuestions = markOcrStudentQuestionService.listByStudentIdAndQuestionId(studentId, markQuestion.getId());
            if (CollectionUtils.isNotEmpty(markOcrStudentQuestions)) {
                String result = markOcrStudentQuestions.stream().map(MarkOcrStudentQuestion::getOcrContent).collect(Collectors.joining(";"));
                dto.setOcrResult(result);
            }
        } catch (ApiException e) {

        } catch (IOException e) {

        } finally {
//            if (CollectionUtils.isNotEmpty(deleteFileList)) {
//                for (File file : deleteFileList) {
//                    FileUtil.deleteFile(file);
//                }
//            }
        }
    }

    private List<MarkTask> listMarkTaskByQuestionId(Long questionId) {
        QueryWrapper<MarkTask> queryWrapper = new QueryWrapper<>();
        queryWrapper.lambda().eq(MarkTask::getQuestionId, questionId)
                .orderByAsc(MarkTask::getStudentId);
        return markTaskService.list(queryWrapper);
    }

    public void export(List<ExportAiResultDto> list) {
        try {
            File file = new File(rootPath + File.separator + "学生名单" + SystemConstant.EXCEL_PREFIX);
            if (!file.exists()) {
                file.getParentFile().mkdirs();
                file.createNewFile();
            }
            FileOutputStream outputStream = new FileOutputStream(file);
            ExcelWriter writer = ExcelWriter.create(ExcelType.XLSX);
            writer.writeObjects("学生名单", null, ExportAiResultDto.class, list.listIterator());
            writer.output(outputStream);
            outputStream.flush();
            outputStream.close();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
}