LaiTool/src/main/Service/subtitle.ts

import { isEmpty } from 'lodash'
import { BookService } from '../../define/db/service/Book/bookService'
import { errorMessage, successMessage } from '../Public/generalTools'
import { FfmpegOptions } from './ffmpegOptions'
import { SubtitleSavePositionType } from '../../define/enum/waterMarkAndSubtitle'
import { BookTaskDetailService } from '../../define/db/service/Book/bookTaskDetailService'
import { define } from '../../define/define'
import path from 'path'
import {
  CheckFileOrDirExist,
  CheckFolderExistsOrCreate,
  DeleteFolderAllFile,
  GetFilesWithExtensions
} from '../../define/Tools/file'
import { shell } from 'electron'
import { Book } from '../../model/book'
import fs from 'fs'
const util = require('util')
const { exec } = require('child_process')
const execAsync = util.promisify(exec)
const fspromises = fs.promises

/**
 * 去除水印和获取字幕相关操作
 */
export class Subtitle {
  bookService: BookService
  bookTaskDetailService: BookTaskDetailService
  ffmpegOptions: FfmpegOptions

  constructor() { }

  async InitService() {
    if (!this.bookService) {
      this.bookService = await BookService.getInstance()
    }
    if (!this.bookTaskDetailService) {
      this.bookTaskDetailService = await BookTaskDetailService.getInstance()
    }
    if (!this.ffmpegOptions) {
      this.ffmpegOptions = new FfmpegOptions()
    }
  }

  //#region 通用方法

  /**
   * 拆分视频总帧数，每秒多少帧，平分视频总帧数，后截取
   * @param {*} videoDurationMs 视频的总时长（毫秒）
   * @param {*} framesPerSecond 每秒截取多少帧
   * @returns
   */
  GenerateFrameTimes(videoDurationMs, framesPerSecond) {
    // 直接使用视频总时长（毫秒），不进行向下取整
    const videoDurationSec = videoDurationMs / 1000

    // 计算总共需要抽取的帧数，考虑到视频时长可能不是完整秒数，使用 Math.ceil 来确保至少获取到最后一秒内的帧
    const totalFrames = Math.ceil(videoDurationSec * framesPerSecond)

    // 计算两帧之间的时间间隔（毫秒）
    const interval = 1000 / framesPerSecond

    // 生成对应的时间点数组
    const frameTimes = []
    for (let i = 0; i < totalFrames; i++) {
      // 使用 Math.min 确保最后一个时间点不会超过视频总时长
      let timePoint = Math.min(Math.round(interval * i), videoDurationMs)
      frameTimes.push(timePoint)
    }

    return frameTimes
  }

  //#endregion

  /**
   * 获取当前视频中所有的字幕信息
   * @param {*} value 需要的参数的对象，包含下面的参数
   * @param {*} value.id 小说ID/小说分镜详细信息ID/null
   * @param {*} value.type 保存的类型（主视频/分镜视频/后续会添加外部单独的视频提取）
   * @param {*} value.videoPath 视频路径
   * @param {*} value.subtitlePosition 字幕位置信息
   */
  async GetVideoFrameText(value: Book.GetVideoFrameTextParams) {
    try {
      await this.InitService()
      let videoPath
      let tempImageFolder
      let position
      if (value.type == SubtitleSavePositionType.MAIN_VIDEO) {
        let bookRes = this.bookService.GetBookDataById(value.id)
        if (bookRes == null) {
          throw new Error('没有找到小说对应的的视频地址')
        }
        let book = bookRes
        tempImageFolder = path.join(define.project_path, `${book.id}/data/subtitle/${book.id}/temp`)
        if (isEmpty(book.subtitlePosition)) {
          throw new Error('请先保存位置信息')
        }
        position = JSON.parse(book.subtitlePosition)
        videoPath = book.oldVideoPath
      } else if (value.type == SubtitleSavePositionType.STORYBOARD_VIDEO) {
        let bookTaskDetail = this.bookTaskDetailService.GetBookTaskDetailDataById(value.id);
        if (bookTaskDetail == null) {
          throw new Error("没有找到小说分镜详细信息")
        }

        tempImageFolder = path.join(define.project_path, `${bookTaskDetail.bookId}/data/subtitle/${bookTaskDetail.name}_${bookTaskDetail.id}/temp`)
        if (isEmpty(value.subtitlePosition)) {
          throw new Error('请先保存位置信息')
        }
        position = JSON.parse(value.subtitlePosition)
        videoPath = bookTaskDetail.videoPath
      } else {
        throw new Error("不支持的操作");
      }

      await CheckFolderExistsOrCreate(tempImageFolder)
      // 判断文件夹是不是存在，存在的话，将里面的所有文件删除
      await DeleteFolderAllFile(tempImageFolder)

      // 将视频进行抽帧，（目前是每秒1帧，时间小于一秒，抽一帧）
      let getDurationRes = await this.ffmpegOptions.FfmpegGetVideoDuration(videoPath)
      if (getDurationRes.code == 0) {
        throw new Error(getDurationRes.message)
      }
      let videoDuration = getDurationRes.data
      let frameTime = this.GenerateFrameTimes(videoDuration, 1)
      for (let i = 0; i < frameTime.length; i++) {
        const item = frameTime[i];
        let name = i.toString().padStart(6, '0')
        let imagePath = path.join(tempImageFolder, `frame_${name}.png`)
        // 开始抽帧
        let res = await this.ffmpegOptions.FfmpegGetVideoFramdAndClip(
          videoPath,
          item,
          imagePath,
          position
        )
        if (res.code == 0) {
          throw new Error(res.message)
        }
      }

      // 开始识别
      let textRes = await this.GetCurrentFrameText({
        id: value.id,
        type: value.type,
        imageFolder: tempImageFolder
      })

      let allTextData = [] as string[]
      // 开始获取所有的数据
      let jsonPaths = await GetFilesWithExtensions(tempImageFolder, ['.json'])
      for (let i = 0; i < jsonPaths.length; i++) {
        const element = jsonPaths[i]
        // 开始拼接
        let texts = JSON.parse(await fspromises.readFile(element, 'utf-8'))
        for (let j = 0; j < texts.length; j++) {
          const text = texts[j][1][0]
          allTextData.includes(text) ? null : allTextData.push(text)
        }
      }

      console.log(allTextData.join('，'))
      // 这边计算相似度，返回过于相似的数据
      // let res = await RemoveSimilarTexts(allTextData)

      return successMessage(
        allTextData.join('，'),
        '获取视频的的文案信息成功',
        'WatermarkAndSubtitle_GetVideoFrameText'
      )
    } catch (error) {
      return errorMessage(
        '提取视频的的文案信息失败，错误消息如下：' + error.toString(),
        'WatermarkAndSubtitle_GetCurrentFrameText'
      )
    }
  }

  /**
   * 获取当前帧的文字信息
   * @param {*} value 需要的参数的对象，必须包含以下参数
   * @param {*} value.id 小说ID/小说分镜详细信息ID/null
   * @param {*} value.type 保存的类型（主视频/分镜视频/后续会添加外部单独的视频提取）
   */
  async GetCurrentFrameText(value) {
    try {
      await this.InitService()
      let iamgePaths = []
      let imageFolder = value.imageFolder
        ? value.imageFolder
        : path.join(define.project_path, `${value.id}/data/subtitle/${value.id}`)
      let imageFolderIsExist = await CheckFileOrDirExist(imageFolder)
      if (!imageFolderIsExist) {
        throw new Error('请先保存位置信息')
      }

      let images = await GetFilesWithExtensions(imageFolder, ['.png'])
      let regex = /.*frame_.*\.png$/
      images.forEach((element) => {
        // 使用正则表达式测试文件名
        if (regex.test(element)) {
          iamgePaths.push(element)
        }
      })

      // 开始识别
      for (let i = 0; i < iamgePaths.length; i++) {
        const imagePath = iamgePaths[i]
        let scriptPath = path.join(define.scripts_path, 'LaiOcr/LaiOcr.exe')
        let script = `cd "${path.dirname(scriptPath)}" && "${scriptPath}" "${imagePath}"`
        let scriptRes = await execAsync(script, { maxBuffer: 1024 * 1024 * 10, encoding: 'utf-8' })
        console.log(scriptRes)
        if (scriptRes.error) {
          throw new Error(scriptRes.error)
        }
      }

      // 处理所有的图片完毕，遍历所有的数据返回
      let textData = []
      let jsonPath = await GetFilesWithExtensions(imageFolder, ['.json'])
      for (let i = 0; i < jsonPath.length; i++) {
        const element = jsonPath[i]
        // 开始拼接
        let texts = JSON.parse(await fspromises.readFile(element, 'utf-8'))
        for (let j = 0; j < texts.length; j++) {
          const text = texts[j][1][0]
          textData.includes(text) ? null : textData.push(text)
        }
      }
      return successMessage(
        textData.join('\n'),
        '获取当前帧的文字信息成功',
        'WatermarkAndSubtitle_GetCurrentFrameText'
      )
    } catch (error) {
      return errorMessage(
        '获取当前帧的文字信息失败，错误消息如下：' + error.toString(),
        'WatermarkAndSubtitle_GetCurrentFrameText'
      )
    }
  }

  /**
   * 打开对应的ID的字幕提取的图片文件夹
   * @param {*} value 需要的参数的对象，必须包含以下参数
   * @param {*} value.id 小说ID/小说分镜详细信息ID/null
   * @param {*} value.type 保存的类型（主视频/分镜视频/后续会添加外部单独的视频提取）
   */
  async OpenBookSubtitlePositionScreenshot(value) {
    try {
      let folder
      if (
        value.type == SubtitleSavePositionType.MAIN_VIDEO ||
        value.type == SubtitleSavePositionType.SETTING
      ) {
        folder = path.join(define.project_path, `${value.id}/data/subtitle/${value.id}`)
      } else if (value.type == SubtitleSavePositionType.STORYBOARD_VIDEO) {
        folder = path.join(define.project_path, `${value.id}/data/subtitle/${value.id}`)
      }

      // 判断文件夹是不是存在
      let folderIsExist = await CheckFileOrDirExist(folder)
      if (!folderIsExist) {
        throw new Error('文件夹不存在，请先保存字幕位置信息')
      }

      // 打开文件夹
      shell.openPath(folder)
      return successMessage(
        null,
        '打开对应的文件夹成功',
        'WatermarkAndSubtitle_OpenBookSubtitlePositionScreenshot'
      )
    } catch (error) {
      return errorMessage(
        '打开字幕位置信息失败，错误消息如下：' + error.toString(),
        'WatermarkAndSubtitle_OpenBookSubtitlePositionScreenshot'
      )
    }
  }

  /**
   * 保存反推的视频的文案位置信息（可以保存多个）
   * @param {*} value 需要的参数的对象，必须包含以下参数
   * @param {*} value.id 小说ID/小说分镜详细信息ID/null
   * @param {*} value.bookSubtitlePosition 小说文案对应的位置
   * @param {*} value.currentTime 视频当前保存的时间
   * @param {*} value.type 保存的类型（主视频/分镜视频/后续会添加外部单独的视频提取）
   * @returns
   */
  async SaveBookSubtitlePosition(value) {
    try {
      await this.InitService()
      let saveData = []
      let videoPath
      let outImagePath
      // 小说视频保存
      this.ffmpegOptions = new FfmpegOptions()
      if (
        value.type == SubtitleSavePositionType.MAIN_VIDEO ||
        value.type == SubtitleSavePositionType.SETTING
      ) {
        if (value.id == null) {
          throw new Error('小说ID不能为空')
        }
        // 获取指定的小说
        let bookRes = this.bookService.GetBookDataById(value.id)
        if (bookRes == null) {
          throw new Error('没有找到小说信息')
        }
        let book = bookRes

        if (value.bookSubtitlePosition.length <= 0) {
          throw new Error('没有获取到字幕信息')
        }

        videoPath = book.oldVideoPath
        if (isEmpty(videoPath)) {
          throw new Error('没有获取到视频路径')
        }
        outImagePath = path.join(book.bookFolderPath, `data/subtitle/${book.id}/frame.png`)

        // 获取视频的宽高数据
        let videoSizeRes = await this.ffmpegOptions.FfmpegGetVideoSize(videoPath)
        if (videoSizeRes.code == 0) {
          throw new Error(videoSizeRes.message)
        }
        let videoSize = videoSizeRes.data

        // 开始计算比例
        let videoWidth = videoSize.width
        let videoHeight = videoSize.height
        for (let i = 0; i < value.bookSubtitlePosition.length; i++) {
          const element = value.bookSubtitlePosition[i]
          let widthRate = videoWidth / element.videoWidth // 宽度比例
          let heightRate = videoHeight / element.videoHeight // 高度比例
          // 计算比例
          let newStartX = widthRate * element.startX
          let newStartY = heightRate * element.startY
          let newWidth = widthRate * element.width
          let newHeight = heightRate * element.height
          saveData.push({
            startX: newStartX,
            startY: newStartY,
            width: newWidth,
            height: newHeight,
            videoWidth: videoWidth,
            videoHeight: videoHeight
          })
        }

        // 数据保存
        let saveRes = await this.bookService.UpdateBookData(value.id, {
          subtitlePosition: JSON.stringify(saveData)
        })
        if (saveRes.code == 0) {
          throw new Error(saveRes.message)
        }
      } else if (value.type == SubtitleSavePositionType.STORYBOARD_VIDEO) {
        // 小说分镜详细信息保存
        if (value.id == null) {
          throw new Error('小说分镜详细信息ID不能为空')
        }
        // 获取指定的小说分镜详细信息
        let bookStoryboard = this.bookTaskDetailService.GetBookTaskDetailDataById(value.id)
        if (bookStoryboard == null) {
          throw new Error('没有找到小说分镜信息')
        }

        if (value.bookSubtitlePosition.length <= 0) {
          throw new Error('没有获取到字幕信息')
        }

        videoPath = bookStoryboard.videoPath
        if (isEmpty(videoPath)) {
          throw new Error('没有获取到视频路径')
        }

        outImagePath = path.join(
          define.project_path,
          `${bookStoryboard.bookId}/data/subtitle/${bookStoryboard.id}/frame.png`
        )

        // 获取视频的宽高数据
        this.ffmpegOptions = new FfmpegOptions()
        let videoSizeRes = await this.ffmpegOptions.FfmpegGetVideoSize(videoPath)
        if (videoSizeRes.code == 0) {
          throw new Error(videoSizeRes.message)
        }
        let videoSize = videoSizeRes.data

        // 开始计算比例
        let videoWidth = videoSize.width
        let videoHeight = videoSize.height
        for (let i = 0; i < value.bookSubtitlePosition.length; i++) {
          const element = value.bookSubtitlePosition[i]
          let widthRate = videoWidth / element.videoWidth // 宽度比例
          let heightRate = videoHeight / element.videoHeight // 高度比例
          // 计算比例
          let newStartX = widthRate * element.startX
          let newStartY = heightRate * element.startY
          let newWidth = widthRate * element.width
          let newHeight = heightRate * element.height
          saveData.push({
            startX: newStartX,
            startY: newStartY,
            width: newWidth,
            height: newHeight,
            videoWidth: videoWidth,
            videoHeight: videoHeight
          })
        }
        // 数据保存
        let saveRes = this.bookTaskDetailService.UpdateBookTaskDetail(bookStoryboard.id, {
          subtitlePosition: JSON.stringify(saveData)
        })
        if (saveRes.code == 0) {
          throw new Error(saveRes.message)
        }
      }

      // 开始设置裁剪出来的图片位置
      // 裁剪一个示例图片
      let saveImagePath = await this.ffmpegOptions.FfmpegGetVideoFramdAndClip(
        videoPath,
        value.currentTime * 1000,
        outImagePath,
        saveData
      )
      if (saveImagePath.code == 0) {
        throw new Error(saveImagePath.message)
      }
      return successMessage(
        saveImagePath.data,
        '保存字幕位置信息成功',
        'WatermarkAndSubtitle_SaveBookSubtitlePosition'
      )
    } catch (error) {
      return errorMessage(
        '保存字幕位置信息失败，错误消息如下：' + error.toString(),
        'WatermarkAndSubtitle_SaveBookSubtitlePosition'
      )
    }
  }
}