shouyoudianjing/src/main/java/com/izouma/awesomeadmin/util/VideoProcessToolNew.java

package com.izouma.awesomeadmin.util;

import com.izouma.awesomeadmin.service.OSSFileService;
import org.apache.commons.lang3.RandomStringUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.time.DateFormatUtils;
import org.bytedeco.javacpp.indexer.UByteRawIndexer;
import org.bytedeco.javacpp.tesseract;
import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.FrameGrabber;
import org.bytedeco.javacv.OpenCVFrameConverter;
import org.hibernate.validator.internal.util.privilegedactions.GetResource;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.regex.Pattern;

import static org.bytedeco.javacpp.lept.pixRead;
import static org.bytedeco.javacpp.opencv_core.*;
import static org.bytedeco.javacpp.opencv_imgcodecs.imencode;
import static org.bytedeco.javacpp.opencv_imgproc.*;
import static org.bytedeco.javacpp.opencv_imgproc.bilateralFilter;
import static org.bytedeco.javacpp.opencv_ml.SVM;
import static org.bytedeco.javacpp.tesseract.PSM_SINGLE_LINE;

public class VideoProcessToolNew {
    private OSSFileService ossFileService = new OSSFileService();
    public  String         imgPrefix      = "/tmp/";

    public Map<String, String> processVideo(String path, int frameSkip) throws FrameGrabber.Exception {
        SVM svm = SVM.load(GetResource.class.getClassLoader().getResource("trainneddata/pubg.xml").getPath());
        Map<String, String> map = new HashMap<>();
        OpenCVFrameConverter.ToMat converterToMat = new OpenCVFrameConverter.ToMat();
        long frameCount = 0;
        FrameGrabber videoGrabber = new FFmpegFrameGrabber(path);
        videoGrabber.start();
        Frame vFrame;
        do {
            vFrame = videoGrabber.grabFrame();
            if (vFrame != null) {
                frameCount++;
                if (frameCount % 200 == 0) {
                    System.gc();
                }
                if (frameSkip > 0 && frameCount % frameSkip != 0) {
                    continue;
                }
                Mat frame = removeBlackBarAndRotate(converterToMat.convert(vFrame));
                if (matchGameOver(svm, frame)) {
                    Mat filtered = new Mat();
                    bilateralFilter(frame, filtered, 25, 25 * 2, 25 / 2f);
                    if (map.get("rank") == null) {
                        String rank = extractRank(filtered);
                        if (StringUtils.isNotEmpty(rank))
                            map.put("rank", rank);
                    }
                    if (map.get("total") == null) {
                        String total = extractTotalNum(filtered);
                        if (StringUtils.isNotEmpty(total))
                            map.put("total", total);
                    }
                    if (map.get("参赛时间") == null) {
                        Map<String, String> statistics = getStatistics(filtered);
                        if (statistics != null)
                            map.putAll(statistics);
                    }
                    filtered.release();
                    if (map.get("rank") != null
                            && map.get("total") != null
                            && map.get("参赛时间") != null
                            && map.get("评分") != null) {
                        double score = 0;
                        double total = 0;
                        double rank = 0;

                        try {
                            score = Double.parseDouble(map.get("评分"));
                            total = Double.parseDouble(map.get("total"));
                            rank = Double.parseDouble(map.get("rank").replace("第", ""));

                            if (score > (total - rank) / total * 100 / 2) {
                                String filename = "/var/samples/" + DateFormatUtils.format(new Date(), "yyyyMMddHHmmss") + RandomStringUtils.randomNumeric(8) + ".jpg";
                                org.bytedeco.javacpp.opencv_imgcodecs.imwrite(filename, frame);
                                map.put("image", uploadImage(frame));
                                System.out.println(map);
                                break;
                            }
                        } catch (Exception ignore) {
                        }
                    }
                }
                frame.release();
                System.out.println(frameCount + " frame processed");
            }
        } while (vFrame != null);
        videoGrabber.stop();
        videoGrabber.release();
        svm.deallocate();
        System.gc();
        return map;
    }

    public boolean matchGameOver(SVM svm, Mat inputImage) {
        Mat resized = new Mat();
        resize(inputImage, resized, new Size(854, 480));
        Mat reshaped = resized.reshape(1, 1);
        Mat toPredict = new Mat();
        reshaped.convertTo(toPredict, CV_32F);
        float res = svm.predict(toPredict);

        resized.release();
        reshaped.release();
        toPredict.release();
        return res == 1;
    }


    public String extractRank(Mat src) {
        String result = null;
        int m = src.rows();
        int n = src.cols();
        Mat roi = new Mat(src, new Range(0, (int) (m * 0.093)), new Range((int) (n * 0.44), (int) (n * 0.56)));
        imwrite("rank.jpg", roi);

        Mat filtered = new Mat(roi.rows(), roi.cols(), CV_8UC1, new Scalar(0));
        Mat hsvImg = new Mat();
        cvtColor(roi, hsvImg, COLOR_BGR2HSV);

        UByteRawIndexer hsvIndexer = hsvImg.createIndexer();
        UByteRawIndexer filteredIndexer = filtered.createIndexer();
        for (int i = 0; i < roi.rows(); i++) {
            for (int j = 0; j < roi.cols(); j++) {
                int h = hsvIndexer.get(i, j, 0);
                int s = hsvIndexer.get(i, j, 1);
                int v = hsvIndexer.get(i, j, 2);
                int dh = 0;
                if (h > 26) {
                    dh = h - 26;
                } else if (h < 21) {
                    dh = 21 - h;
                }
                int ds = 0;
                if (s > 161) {
                    ds = s - 161;
                } else if (s < 120) {
                    ds = 120 - s;
                }
                int dv = 0;
                if (v > 251) {
                    dv = v - 251;
                } else if (v < 238) {
                    dv = 238 - v;
                }
                int gray = (int) (255 - Math.min(255, dh * 50)
                        - Math.min(255, ds * 3)
                        - Math.min(255, dv * 10));
                gray = gray < 0 ? 0 : gray;
                // int gray = (h == 26 && Math.abs(s - S) < 50) ? 255 : 0;
                filteredIndexer.put(i, j, gray);
            }
        }
        hsvIndexer.release();
        hsvImg.release();
        filteredIndexer.release();

        imwrite("rank_filtered.jpg", filtered);
        Mat gray = new Mat();
        adaptiveThreshold(filtered, gray, 255, CV_ADAPTIVE_THRESH_MEAN_C, CV_THRESH_BINARY_INV, 15, 12);
        imwrite("rank_gray.jpg", gray);
        MatVector contours = new MatVector();
        findContours(gray, contours, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);
        List<Rect> rects = new ArrayList<>();
        for (int i = 0; i < contours.size(); i++) {
            Mat contour = contours.get(i);
            Rect rect = boundingRect(contour);
            rects.add(rect);
            contour.release();
        }
        rects.sort((o1, o2) -> o2.area() - o1.area());
        for (int i = 0; i < rects.size(); i++) {
            Rect rect = rects.get(i);
            if (i < 5) {
                Mat txtImg = new Mat(filtered, new Range(rect.y(), rect.y() + rect.height()), new Range(rect.x(), rect.x() + rect.width()));
                String str = doOCR(txtImg);
                txtImg.release();
                System.out.println(str);
                if (str != null) {
                    str = str.replaceAll("[ \n]", "");
                    if (Pattern.matches("^第\\d{1,3}$", str)) {
                        result = str;
                        break;
                    }
                }
            }
        }
        filtered.release();
        gray.release();
        roi.release();
        return result;
    }

    public String extractTotalNum(Mat src) {
        String result = null;
        int m = src.rows();
        int n = src.cols();
        Mat roi = new Mat(src, new Range(0, (int) (m * 0.08)), new Range((int) (n * 0.42), (int) (n * 0.58)));
        imwrite("total.png", roi);

        Mat totalFiltered = new Mat(roi.rows(), roi.cols(), CV_8UC1, new Scalar(0));
        Mat hsvImg = new Mat();
        cvtColor(roi, hsvImg, COLOR_BGR2HSV);

        UByteRawIndexer hsvIndexer = hsvImg.createIndexer();
        UByteRawIndexer resultIndexer = totalFiltered.createIndexer();
        for (int i = 0; i < roi.rows(); i++) {
            for (int j = 0; j < roi.cols(); j++) {
                int h = hsvIndexer.get(i, j, 0);
                int s = hsvIndexer.get(i, j, 1);
                int v = hsvIndexer.get(i, j, 2);

                int gray = 255 - 2 * s;
                gray = gray < 0 ? 0 : gray;
                resultIndexer.put(i, j, gray);
            }
        }
        hsvIndexer.release();
        hsvImg.release();
        resultIndexer.release();
        imwrite("total_filtered.png", totalFiltered);
        Mat gray = new Mat();
        adaptiveThreshold(totalFiltered, gray, 255, CV_ADAPTIVE_THRESH_MEAN_C, CV_THRESH_BINARY_INV, 15, 12);
        imwrite("rank_gray.jpg", gray);
        MatVector contours = new MatVector();
        findContours(gray, contours, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);
        List<Rect> rects = new ArrayList<>();
        for (int i = 0; i < contours.size(); i++) {
            Mat contour = contours.get(i);
            Rect rect = boundingRect(contour);
            rects.add(rect);
            contour.release();
        }
        rects.sort((o1, o2) -> o2.area() - o1.area());
        for (int i = 0; i < rects.size(); i++) {
            Rect rect = rects.get(i);
            if (i < 5) {
                Mat txtImg = new Mat(totalFiltered, new Range(rect.y(), rect.y() + rect.height()), new Range(rect.x(), rect.x() + rect.width()));
                String str = doOCR(txtImg);
                txtImg.release();
                if (str != null) {
                    str = str.replaceAll("[ \n]", "");
                    if (Pattern.matches("^/\\d{1,3}$", str)) {
                        result = str.replace("/", "");
                        break;
                    }
                }
            }
        }
        totalFiltered.release();
        gray.release();
        System.out.println(result);
        return result;
    }

    public String extractMode(Mat src) {
        String result = null;
        int m = src.rows();
        int n = src.cols();
        Mat roi = new Mat(src, new Range(0, (int) (m * 0.065)), new Range(0, (int) (n * 0.2)));
        imwrite("mode.png", roi);

        Mat filtered = new Mat(roi.rows(), roi.cols(), CV_8UC1, new Scalar(0));
        Mat hsvImg = new Mat();
        cvtColor(roi, hsvImg, COLOR_BGR2HSV);

        UByteRawIndexer hsvIndexer = hsvImg.createIndexer();
        UByteRawIndexer bgrIndexer = roi.createIndexer();
        UByteRawIndexer resultIndexer = filtered.createIndexer();
        for (int i = 0; i < roi.rows(); i++) {
            for (int j = 0; j < roi.cols(); j++) {
                int h = hsvIndexer.get(i, j, 0);
                int s = hsvIndexer.get(i, j, 1);
                int v = hsvIndexer.get(i, j, 2);
                int b = bgrIndexer.get(i, j, 0);
                int g = bgrIndexer.get(i, j, 1);
                int r = bgrIndexer.get(i, j, 2);

                int dr = Math.abs(r - 220);
                int dg = Math.abs(g - 220);
                int db = Math.abs(b - 220);

                int dbgr = Math.abs(b - g) + Math.abs(b - r) + Math.abs(r - g);

                int gray = (int) (255 - Math.pow(dr, 2) * 0.015 - Math.pow(dg, 2) * 0.015 - Math.pow(db, 2) * 0.015 - Math.pow(dbgr, 2) * 0.015);
                // gray = (int) (gray / 128f * 255);
                // gray = 255 - Math.abs(gray - 200) * 2;
                if (gray > 220) {
                    gray = 0;
                } else if (gray > 190) {
                    gray = 220 - 190;
                } else {
                    gray = 255;
                }
                resultIndexer.put(i, j, gray);
            }
        }

        hsvIndexer.release();
        bgrIndexer.release();
        hsvImg.release();
        resultIndexer.release();
        imwrite("mode_filtered.png", filtered);
        result = doOCR(filtered, "chi_sim");
        if (StringUtils.isNotEmpty(result)) {
            result = result.replace(" ", "").replace("\n", "");
        }
        System.out.println(result);
        filtered.release();
        roi.release();
        return result;
    }

    public Map<String, String> getStatistics(Mat src) {
        imwrite("time.jpg", src);
        int m = src.rows();
        int n = src.cols();
        Mat roi = new Mat(src, new Range((int) (m * 0.569), (int) (m * 0.625)), new Range((int) (n * 0.3), (int) (n * 0.8)));
        Mat filtered = timeFilter(roi);
        imwrite("time_filtered.jpg", filtered);

        Map<String, String> map = new HashMap<>();
        Rect[] rects = {
                new Rect(0, 0, (int) (filtered.cols() * 0.128), filtered.rows()),
                // new Rect((int) (filtered.cols() * 0.128), 0, (int) (filtered.cols() * 0.128), filtered.rows()),
                new Rect((int) (filtered.cols() * 0.256), 0, (int) (filtered.cols() * 0.162), filtered.rows()),
                // new Rect((int) (filtered.cols() * 0.428), 0, (int) (filtered.cols() * 0.114), filtered.rows()),
                // new Rect((int) (filtered.cols() * 0.542), 0, (int) (filtered.cols() * 0.1), filtered.rows()),
                // new Rect((int) (filtered.cols() * 0.642), 0, (int) (filtered.cols() * 0.1), filtered.rows()),
                new Rect((int) (filtered.cols() * 0.742), 0, (int) (filtered.cols() * 0.1), filtered.rows())
        };
        String[] names = {
                "淘汰",
                // "伤害",
                "参赛时间",
                // "治疗量",
                // "救援",
                // "助攻",
                "评分"
        };
        int i = 0;
        for (Rect rect : rects) {
            Mat txtImg = new Mat(filtered, new Range(rect.y(), rect.y() + rect.height()), new Range(rect.x(), rect.x() + rect.width()));

            Rect rect1 = boundingRect(txtImg);
            int y1 = rect1.y() - 5;
            y1 = y1 < 0 ? 0 : y1;
            int y2 = rect1.y() + rect1.height() + 5;
            y2 = y2 > (txtImg.rows() - 1) ? (txtImg.rows() - 1) : y2;
            int x1 = rect1.x() - 5;
            x1 = x1 < 0 ? 0 : x1;
            int x2 = rect1.x() + rect1.width() + 5;
            x2 = x2 > (txtImg.cols() - 1) ? (txtImg.cols() - 1) : x2;
            txtImg = txtImg.rowRange(y1, y2).colRange(x1, x2);
            bitwise_not(txtImg, txtImg);
            resize(txtImg, txtImg, new Size(0, 0), 3, 3, INTER_LINEAR);
            String str = doOCR(txtImg);
            if (StringUtils.isNotEmpty(str)) {
                str = str.replaceAll("[ \n]", "");
            }
            map.put(names[i], str);
            System.out.println(names[i] + str);
            imwrite(names[i] + ".png", txtImg);
            i++;
        }

        filtered.release();
        roi.release();
        if (Pattern.matches("((^([0-9]+[.][0-9]*))|(^\\d{1,2}))分钟", map.get("参赛时间")) &&
                Pattern.matches("((^([0-9]+[.][0-9]+)$)|(^\\d{1,2})$)", map.get("评分"))) {
            try {
                double score = Double.parseDouble(map.get("评分"));
                double time = Double.parseDouble(map.get("参赛时间").replace("分钟", ""));
                if (score >= 0 && score <= 100 && time >= 0 && time <= 45) {
                    return map;
                }
            } catch (Exception ignored) {
            }
        }
        return null;
    }

    public Mat timeFilter(Mat src) {
        int b = 27;
        int g = 145;
        int r = 211;

        int hmin = 28;
        int hmax = 28;
        int smin = 145;
        int smax = 145;
        int vmin = 160;
        int vmax = 160;

        Mat result = new Mat(src.rows(), src.cols(), CV_8UC1, new Scalar(0));
        Mat hsvImg = new Mat();
        cvtColor(src, hsvImg, COLOR_BGR2HSV);
        UByteRawIndexer hsvIndexer = hsvImg.createIndexer();
        UByteRawIndexer srcIndexer = src.createIndexer();
        UByteRawIndexer resultIndexer = result.createIndexer();
        for (int i = 0; i < src.rows(); i++) {
            for (int j = 0; j < src.cols(); j++) {
                int db = Math.abs(srcIndexer.get(i, j, 0) - b);
                int dg = Math.abs(srcIndexer.get(i, j, 1) - g);
                int dr = Math.abs(srcIndexer.get(i, j, 2) - r);

                int h = hsvIndexer.get(i, j, 0);
                int s = hsvIndexer.get(i, j, 1);
                int v = hsvIndexer.get(i, j, 2);

                int dh = h > hmax ? (h - hmax) : (h < hmin ? (hmin - h) : 0);
                int ds = s > smax ? (s - smax) : (s < smin ? (smin - s) : 0);
                int dv = v > vmax ? (v - vmax) : (v < vmin ? (vmin - v) : 0);

                int gray = 255 - (db * 0 + dg * 3 + dr * 2);
                // int gray = 255 - dh * 2 - ds * 2 - dv* 2;
                if (gray < 0) gray = 0;
                if (gray > 255) gray = 255;
                gray *= 1.2;
                if (gray > 255) gray = 255;
                resultIndexer.put(i, j, gray);
            }
        }
        hsvImg.release();
        hsvIndexer.release();
        srcIndexer.release();
        resultIndexer.release();
        return result;
    }


    private String doOCR(Mat img) {
        return doOCR(img, "pubg");
    }

    private String doOCR(Mat img, String lang) {
        imwrite("text.jpg", img);
        final tesseract.TessBaseAPI baseApi = new tesseract.TessBaseAPI();
        baseApi.Init(GetResource.class.getClassLoader().getResource("trainneddata").getPath(), lang);
        baseApi.SetPageSegMode(PSM_SINGLE_LINE);
        baseApi.SetImage(pixRead(imgPrefix + "text.jpg"));
        String recognizedText = baseApi.GetUTF8Text().getString();
        return recognizedText == null ? "" : recognizedText;
    }

    public Mat removeBlackBarAndRotate(Mat src) {
        if (src.cols() < src.rows()) {
            rotate(src, src, ROTATE_90_COUNTERCLOCKWISE);
        }
        Mat gray = new Mat();
        cvtColor(src, gray, COLOR_BGR2GRAY);
        threshold(gray, gray, 2, 255, THRESH_BINARY);
        Rect rect = boundingRect(gray);
        Mat dst = new Mat(src, new Range(rect.y(), rect.y() + rect.height()), new Range(rect.x(), rect.x() + rect.width()));
        src.release();
        gray.release();
        rect.deallocate();
        return dst;
    }

    public BufferedImage createBufferedImage(Mat mat) {
        ByteBuffer byteBuffer = ByteBuffer.allocate(mat.cols() * mat.rows());
        imencode(".jpg", mat, byteBuffer);
        byte[] ba = byteBuffer.array();
        BufferedImage bi = null;
        try {
            bi = ImageIO.read(new ByteArrayInputStream(ba));
        } catch (IOException e) {
            e.printStackTrace();
        }
        return bi;
    }

    public String uploadImage(Mat frame) {
        try {
            BufferedImage image = createBufferedImage(frame);
            ByteArrayOutputStream baos = new ByteArrayOutputStream();//io流
            ImageIO.write(image, "png", baos);//写入流中
            byte[] bytes = baos.toByteArray();//转换成字节
            InputStream in = new ByteArrayInputStream(bytes);
            String path = String.format("images/%s-%s.jpg", new SimpleDateFormat("yyyy-MM-dd-hh-mm-ss").format(new Date()), UUID.randomUUID().toString());
            return ossFileService.upload(in, path);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    private void imwrite(String name, Mat img) {
        org.bytedeco.javacpp.opencv_imgcodecs.imwrite(imgPrefix + name, img);
    }
}