# setting.py 文件
# 配置文件

video_path = "E:\\kid.mp4"

video_width = 1280
video_high = 720
# 字幕的帧数
video_frame = 30

# 字幕的区域 从上往下  从左往右
subtitle_x = 600
subtitle_y = 0
# 字幕的高度
subtitle_height = video_high-subtitle_x
# 保存图片的位置
subtitle_path = "D://subtitle.png"
# easyOcr 识别会存在一定误差所以需要判断识别的相似度的范围
subtitle_similarity = 0.5


# 窃取的时间范围 单位秒
clip_time_range = [0, 100]

# 扫描的步长  单位帧
steep = 15

#!/usr/bin/env python
# coding: utf-8
# op.py

import easyocr
import setting
import os
import cv2
import difflib
from PIL import Image

os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"

def image_orc(img):
    reader = easyocr.Reader(['ch_sim'])
    #  image: file path or numpy-array or a byte stream object
    return reader.readtext(img)


def read_video(video):
    # 命中帧
    key_subtitle = []
    # 当前帧
    current_frame = 0
    last_subtitle = ''
    while True:
        ret, frame = video.read()
        if ret:
            current_frame += 1
            if current_frame % setting.steep == 0:
                subtitle_area = frame[setting.subtitle_x:, setting.subtitle_y:, :]
                ocr_result = image_orc(subtitle_area);
                if ocr_result is None:
                    continue
                ocr_subtitle = ocr_result[0][1]
                # 相似度
                similarity = difflib.SequenceMatcher(None, last_subtitle, ocr_subtitle).quick_ratio()
                # 关键帧+1 字幕区域保存 字幕暂存
                if similarity < setting.subtitle_similarity:
                    last_subtitle = ocr_subtitle
                    print(ocr_subtitle)
                    key_subtitle.append(subtitle_area)
        else:
            break
    return key_subtitle


if __name__ == '__main__':
    video = cv2.VideoCapture(setting.video_path)
    subtitles = read_video(video)
    subtitle_image = Image.new('RGB', (setting.video_width, setting.subtitle_height * (len(subtitles))))
    for index, item in enumerate(subtitles):
        img = Image.fromarray(item)
        subtitle_image.paste(img,
                             box=(0, index * setting.subtitle_height, setting.video_width,
                                  setting.subtitle_height * (index + 1)))
    subtitle_image.save(setting.subtitle_path)
    subtitle_image.show()