package com.luooqi.ocr;

import cn.hutool.core.lang.UUID;
import com.luooqi.ocr.utils.CommUtils;
import com.luooqi.ocr.utils.OcrUtils;
import org.apache.commons.io.FileUtils;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.Collection;

/**
 * 主窗体
 *
 * @author liuyuqi
 */
public class MainFm {
    private static String outputDirPath;

    public static void main(String[] args) {
        String srcDirPath = "D:/liuyuqi/fishsource/python/video-subtitle-extract/videos/大象解说《血战钢锯岭》";
        outputDirPath = "D:/liuyuqi/fishsource/python/video-subtitle-extract/res";

        // 获取所有png文件
        Collection<File> javaGbkFileCol = FileUtils.listFiles(new File(srcDirPath), new String[]{"png"}, true);
        String res = "";
        for (File img : javaGbkFileCol) {
//            对每个图片识别
            res += imgOCR(img) + "\r\n";
        }
        saveRes(res);
        System.out.println(res);
    }

    public static void saveRes(String res) {
        File file = new File(outputDirPath + Math.abs(UUID.randomUUID().hashCode()) + ".txt");
        try {
            FileUtils.writeStringToFile(file, res);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public static String imgOCR(File img) {
        BufferedImage image;
        try {
            image = ImageIO.read(img);
        } catch (IOException e) {
            e.printStackTrace();
            return "";
        }
        byte[] bytes = CommUtils.imageToBytes(image);
        int ocrType = 0;
        String text = OcrUtils.ocrImg(bytes, Math.abs(UUID.randomUUID().hashCode()) % 4);
        if (text == "") {
            text = OcrUtils.ocrImg(bytes, Math.abs(UUID.randomUUID().hashCode()) % 4);
        }
        return text;
    }
}