LaiTool/src/main/Service/Book/ReverseBook.ts
2024-08-03 12:46:12 +08:00

589 lines
22 KiB
TypeScript
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import { successMessage, errorMessage } from '../../Public/generalTools'
import { BookBasic } from './BooKBasic'
import { BookService } from '../../../define/db/service/Book/bookService'
import { BookTaskService } from '../../../define/db/service/Book/bookTaskService'
import { define } from '../../../define/define.js'
import fs from 'fs'
import { DEFINE_STRING } from "../../../define/define_string";
import path from 'path'
import { BasicReverse } from './basicReverse'
import { BookTaskDetailService } from '../../../define/db/service/Book/bookTaskDetailService'
import { TaskScheduler } from "../../Service/taskScheduler"
import { Book } from '../../../model/book'
import { LoggerStatus, OtherData, ResponseMessageType } from '../../../define/enum/softwareEnum'
import { GeneralResponse } from '../../../model/generalResponse'
import { cloneDeep, isEmpty } from 'lodash'
import { Subtitle } from '../../Service/subtitle'
import { Watermark } from '../watermark'
import { SubtitleSavePositionType } from '../../../define/enum/waterMarkAndSubtitle'
import { CheckFileOrDirExist, CheckFolderExistsOrCreate, CopyFileOrFolder, GetFilesWithExtensions } from '../../../define/Tools/file'
import { ValidateJson } from '../../../define/Tools/validate'
import { GetImageBase64, ProcessImage } from '../../../define/Tools/image'
import { BookBackTaskType, BookType, TagDefineType, TaskExecuteType } from '../../../define/enum/bookEnum'
import { BookBackTaskListService } from '../../../define/db/service/Book/bookBackTaskListService'
import { MJOpt } from '../MJ/mj'
import { TagDefine } from '../../../define/tagDefine'
import { ImageStyleDefine } from '../../../define/iamgeStyleDefine'
/**
* 一键反推的相关操作
*/
export class ReverseBook extends BookBasic {
basicReverse: BasicReverse
bookTaskService: BookTaskService
taskScheduler: TaskScheduler
bookService: BookService
bookTaskDetailService: BookTaskDetailService
bookBackTaskListService: BookBackTaskListService
mjOpt: MJOpt = new MJOpt()
tagDefine: TagDefine
subtitle: Subtitle
watermark: Watermark
constructor() {
super()
this.basicReverse = new BasicReverse()
this.tagDefine = new TagDefine()
}
async InitService() {
if (!this.bookTaskService) {
this.bookTaskService = await BookTaskService.getInstance()
}
if (!this.taskScheduler) {
this.taskScheduler = new TaskScheduler()
}
if (!this.bookService) {
this.bookService = await BookService.getInstance()
}
if (!this.bookTaskDetailService) {
this.bookTaskDetailService = await BookTaskDetailService.getInstance()
}
if (!this.subtitle) {
this.subtitle = new Subtitle()
}
if (!this.watermark) {
this.watermark = new Watermark()
}
if (!this.bookBackTaskListService) {
this.bookBackTaskListService = await BookBackTaskListService.getInstance()
}
}
// 主动返回前端的消息
sendReturnMessage(data: GeneralResponse.MessageResponse, message_name = DEFINE_STRING.BOOK.GET_COPYWRITING_RETURN) {
global.newWindow[0].win.webContents.send(message_name, data)
}
//#region 小说相关操作
/**
* 获取当前的小说数据
* @param {*} bookQuery
*/
async GetBookData(bookQuery) {
try {
let _bookService = await BookService.getInstance()
// 添加小说
let res = _bookService.GetBookData(bookQuery)
if (res.code == 0) {
throw new Error(res.message)
}
return res
} catch (error) {
return errorMessage(error.message, 'ReverseBook_GetBookData')
}
}
//#endregion
//#region 小说批次任务相关操作
/**
* 获取小说的任务列表
* @param {*} bookTaskCondition 查询任务列表的条件
*/
async GetBookTaskData(bookTaskCondition): Promise<GeneralResponse.ErrorItem | GeneralResponse.SuccessItem> {
try {
await this.InitService()
let _bookTaskService = await BookTaskService.getInstance()
let res = _bookTaskService.GetBookTaskData(bookTaskCondition)
if (res.code == 0) {
throw new Error(res.message)
}
// //TODO 这个后面是不是将所有的数据都是用数据库
// // 这边加载自定义风格
// let styleLists = await this.tagDefine.getTagDataByTypeAndProperty('dynamic', "style_tags");
// if (styleLists.code == 0) {
// throw new Error('获取自定义风格失败')
// }
// let styleTags = styleLists.data as Book.BookStyle[];
// // 异步操作获取风格信息
// for (let index = 0; index < res.data.bookTasks.length; index++) {
// const element = res.data.bookTasks[index];
// let styleList = [] as Book.BookStyle[] | Book.DefineBookStyle[];
// // if (element.imageStyle.length > 0) { // 软件自带的风格
// // let infoRes = ImageStyleDefine.getImageStyleInfomation(JSON.stringify(element.imageStyle));
// // if (infoRes.code == 0) {
// // throw new Error('获取风格失败');
// // }
// // styleList = infoRes.data;
// // }
// if (element.customizeImageStyle.length > 0) { // 自定义的风格
// element.customizeImageStyle.forEach((item) => {
// let style = styleTags.find((tag) => tag.key == item);
// if (style) {
// styleList.push(style);
// }
// });
// }
// res.data.bookTasks[index].styleList = styleList;
// }
return successMessage(res.data, '获取小说任务成功', 'ReverseBook_GetBookTaskData')
} catch (error) {
return errorMessage(
'获取小说对应批次错误,错误信息入校:' + error.message,
'ReverseBook_GetBookTaskData'
)
}
}
/**
* 获取小说的所有的任务详情
* @param bookTaskId
*/
async GetBookTaskDetail(bookTaskId: string) {
try {
await this.InitService()
let _bookTaskDetailService = await BookTaskDetailService.getInstance()
let res = _bookTaskDetailService.GetBookTaskData({ bookTaskId: bookTaskId })
if (res.code == 0) {
throw new Error(res.message)
}
return res
} catch (error) {
return errorMessage(
'获取小说对应批次错误,错误信息入校:' + error.message,
'ReverseBook_GetBookTaskData'
)
}
}
//#endregion
//#region 一键全自动
/**
* 全自动任务(这边是任务入口,都是在这边调用)
* @param {*} value
* @returns
*/
async AutoAction(bookId) {
try {
// 在一键全自动之前当前小说对应的批次任务中的所有的子任务都改为fail然后再执行
// 获取对应的小说小说数据,找到对应的小说视频地址
// let _bookService = await BookService.getInstance()
// let _bookTaskService = await BookTaskService.getInstance()
// let bookData = _bookService.GetBookDataById(bookId)
// if (bookData.data == null) {
// throw new Error('没有找到对应的小说数据请检查bookId是否正确')
// }
// // 获取小说对应的批次任务数据,默认初始化为第一个
// let bookTaskRes = _bookTaskService.GetBookTaskData({
// bookId: bookId,
// name: 'output_00001'
// })
// if (bookTaskRes.data.bookTasks.length <= 0 || bookTaskRes.data.total <= 0) {
// throw new Error('没有找到对应的小说批次任务数据请检查bookId是否正确')
// }
// // 获取小说的视频地址
// let book = bookData.data
// let bookTask = bookTaskRes.data.bookTasks[0]
// // 将当前小说对应的批次任务中的所有的子任务都改为fail
// let updateTaskRes = _bookTaskService.UpdetedBookTaskToFail(bookId, bookTask.id)
// // 添加分镜任务 后面就会全自动的开始执行
// let res = await this.basicReverse.AddFrameDataTask(bookId)
// 添加分割视频任务
// let res = await this.basicReverse.AddCutVideoDataTask(bookId)
// 添加音频分离任务
// let res = await this.basicReverse.AddSplitAudioDataTask(bookId)
// 添加图片抽帧任务
let res = await this.basicReverse.AddGetFrameTask(bookId)
if (res.code == 0) {
throw new Error(res.message)
}
} catch (error) {
return errorMessage(error.message, 'ReverseBook_AutoAction')
}
}
async GetBookAndTask(bookId: string, bookTaskName: string) {
let book = this.bookService.GetBookDataById(bookId)
if (book == null) {
throw new Error("查找小说数据失败");
}
// 获取小说对应的批次任务数据,默认初始化为第一个
let bookTaskRes = await this.bookTaskService.GetBookTaskData({
bookId: bookId,
name: bookTaskName
})
if (bookTaskRes.data.bookTasks.length <= 0 || bookTaskRes.data.total <= 0) {
let msg = "没有找到对应的小说批次任务数据"
this.taskScheduler.AddLogToDB(bookId, book.type, msg, OtherData.DEFAULT, LoggerStatus.FAIL)
throw new Error(msg)
}
return { book: book as Book.SelectBook, bookTask: bookTaskRes.data.bookTasks[0] as Book.SelectBookTask }
}
/**
* 开始分镜任务
*/
async ComputeStoryboard(bookId: string): Promise<any> {
try {
await this.InitService()
let { book, bookTask } = await this.GetBookAndTask(bookId, 'output_00001')
let res = await this.basicReverse.ComputeStoryboardFunc(bookId, bookTask.id);
// 分镜成功直接返回
return successMessage(null, res, "ReverseBook_ComputeStoryboard")
} catch (error) {
return errorMessage("分镜计算失败,失败信息如下 " + error.message, 'ReverseBook_ComputeStoryboard')
}
}
/**
* 开始进行分镜,切割视频并抽帧
* @param bookId
*/
async Framing(bookId: string): Promise<GeneralResponse.ErrorItem | GeneralResponse.SuccessItem> {
try {
await this.InitService()
let { book, bookTask } = await this.GetBookAndTask(bookId, 'output_00001')
// 获取所有的分镜数据
let bookTaskDetail = this.bookTaskDetailService.GetBookTaskData({
bookId: bookId,
bookTaskId: bookTask.id
})
if (bookTaskDetail.data.length <= 0) {
// 传入的分镜数据为空,需要重新获取
await this.taskScheduler.AddLogToDB(
bookId,
book.type,
`没有传入分镜数据,请先进行分镜计算`,
OtherData.DEFAULT,
LoggerStatus.DOING
)
throw new Error("没有传入分镜数据,请先进行分镜计算");
}
let bookTaskDetails = bookTaskDetail.data as Book.SelectBookTaskDetail[]
for (let i = 0; i < bookTaskDetails.length; i++) {
const item = bookTaskDetails[i];
let res = await this.basicReverse.FrameDataToCutVideoData(item);
}
await this.taskScheduler.AddLogToDB(
bookId,
book.type,
"所有的视频裁剪完成,开始抽帧",
bookTask.id,
LoggerStatus.SUCCESS
)
bookTaskDetail = this.bookTaskDetailService.GetBookTaskData({
bookId: bookId,
bookTaskId: bookTask.id
})
bookTaskDetails = bookTaskDetail.data as Book.SelectBookTaskDetail[]
for (let i = 0; i < bookTaskDetails.length; i++) {
const item = bookTaskDetails[i];
let res = await this.basicReverse.FrameFunc(book, item);
}
// 分镜成功直接返回
return successMessage(null, "所有的视频裁剪,抽帧完成", "ReverseBook_Framing")
} catch (error) {
return errorMessage("开始切割视频并抽帧失败,失败信息如下:" + error.message, 'ReverseBook_Framing')
}
}
/**
* 开始执行分镜任务
*/
async GetCopywriting(bookId: string): Promise<GeneralResponse.ErrorItem | GeneralResponse.SuccessItem> {
try {
await this.InitService()
let { book, bookTask } = await this.GetBookAndTask(bookId, 'output_00001')
if (isEmpty(book.subtitlePosition)) {
throw new Error("请先设置小说的字幕位置")
}
// 获取所有的分镜数据
let bookTaskDetail = this.bookTaskDetailService.GetBookTaskData({
bookId: bookId,
bookTaskId: bookTask.id
})
if (bookTaskDetail.data.length <= 0) {
// 传入的分镜数据为空,需要重新获取
await this.taskScheduler.AddLogToDB(
bookId,
book.type,
`没有传入分镜数据,请先进行分镜计算`,
OtherData.DEFAULT,
LoggerStatus.DOING
)
throw new Error("没有传入分镜数据,请先进行分镜计算");
}
let bookTaskDetails = bookTaskDetail.data as Book.SelectBookTaskDetail[]
for (let i = 0; i < bookTaskDetails.length; i++) {
const item = bookTaskDetails[i];
let res = await this.subtitle.GetVideoFrameText({
id: item.id,
videoPath: item.videoPath,
type: SubtitleSavePositionType.STORYBOARD_VIDEO,
subtitlePosition: book.subtitlePosition
})
if (res.code == 0) {
throw new Error(res.message)
}
// 修改数据
this.bookTaskDetailService.UpdateBookTaskDetail(item.id, {
word: res.data,
afterGpt: res.data
});
// let res = await this.basicReverse.GetCopywritingFunc(book, item);
// 将当前的数据实时返回,前端进行修改
this.sendReturnMessage({
code: 1,
id: item.id,
type: ResponseMessageType.GET_TEXT,
data: res.data // 返回识别到的文案
})
// 添加日志
await this.taskScheduler.AddLogToDB(
bookId,
book.type,
`${item.name} 识别文案成功`,
bookTask.id,
LoggerStatus.SUCCESS
)
}
return successMessage(null, "识别是所有文案成功", "ReverseBook_GetCopywriting")
} catch (error) {
return errorMessage("获取分镜数据失败,失败信息如下:" + error.message, 'ReverseBook_GetCopywriting')
}
}
/**
* 执行去除水印操作
* @param bookId 小说ID
*/
async RemoveWatermark(bookId) {
try {
await this.InitService()
let { book, bookTask } = await this.GetBookAndTask(bookId, 'output_00001')
let subtitlePosition = book.subtitlePosition;
if (isEmpty(subtitlePosition)) {
throw new Error("当前没有文案位置的内容,请先进行 ‘开始提取文案 -> 提取文案设置 -> 框选大概位置 -> 保存位置’ 操作保存文案的位置")
}
if (!ValidateJson(subtitlePosition)) {
throw new Error("文案位置格式有误,请先进行 ‘开始提取文案 -> 提取文案设置 -> 框选大概位置 -> 保存位置’ 操作保存文案的位置")
}
let subtitlePositionParse = JSON.parse(subtitlePosition)
if (subtitlePositionParse.length <= 0) {
throw new Error("文案位置格式有误,请先进行 ‘开始提取文案 -> 提取文案设置 -> 框选大概位置 -> 保存位置’ 操作保存文案的位置")
}
let inputImageFolder = path.resolve(define.project_path, `${book.id}/tmp/input`)
if (!(await CheckFileOrDirExist(inputImageFolder))) {
throw new Error("输出文件夹不存在,请先进行抽帧等操作")
}
let iamgePaths = await GetFilesWithExtensions(inputImageFolder, ['.png'])
if (iamgePaths.length <= 0) {
throw new Error("没有检查到抽帧图片,请先进行抽帧等操作")
}
// 处理位置生成蒙板
let inputImage = iamgePaths[0]
let outImagePath = path.resolve(define.project_path, `${book.id}/data/mask/mask_temp_${new Date().getTime()}.png`)
await CheckFolderExistsOrCreate(path.dirname(outImagePath));
await ProcessImage(inputImage, outImagePath, {
x: subtitlePositionParse[0].startX,
y: subtitlePositionParse[0].startY,
width: subtitlePositionParse[0].width,
height: subtitlePositionParse[0].height
})
if (!(await CheckFileOrDirExist(outImagePath))) {
throw new Error("生成蒙板失败")
}
// 开始去除水印
let bookTaskDetails = this.bookTaskDetailService.GetBookTaskData({
bookId: bookId,
bookTaskId: bookTask.id
}).data as Book.SelectBookTaskDetail[]
for (let i = 0; i < bookTaskDetails.length; i++) {
const element = bookTaskDetails[i];
let bookInputImage = element.oldImage;
let name = path.basename(bookInputImage)
let bak = path.resolve(define.project_path, `${book.id}/tmp/input/bak/${name}`);
// 做个备份吧
await CopyFileOrFolder(bookInputImage, bak)
await fs.promises.unlink(bookInputImage) // 删除原来的图片
// 开始调用去除水印的方法
let imageBase64 = await GetImageBase64(bak)
let maskBase64 = await GetImageBase64(outImagePath)
let res = await this.watermark.ProcessImage({
imageBase64: imageBase64,
maskBase64: maskBase64,
type: 'file',
inputFilePath: bak,
maskPath: outImagePath,
outFilePath: bookInputImage
})
// 去水印执行完毕
this.sendReturnMessage({
code: 1,
id: element.id,
type: ResponseMessageType.REMOVE_WATERMARK,
data: res
}, DEFINE_STRING.BOOK.REMOVE_WATERMARK_RETURN)
this.taskScheduler.AddLogToDB(bookId, book.type, `${element.name} 去除水印完成`, element.bookTaskId, LoggerStatus.SUCCESS)
}
// 全部完毕
return successMessage(null, "全部图片去除水印完成", "ReverseBook_RemoveWatermark")
} catch (error) {
return errorMessage("去除水印失败,错误信息如下:" + error.message, "ReverseBook_RemoveWatermark")
}
}
//#endregion
//#region 反推相关任务
/**
* 添加单句生图任务
* @param bookTaskDetailId 小说单个分镜的ID
* @param type 反推的类型
* @returns
*/
async AddReversePromptTask(bookTaskDetailIds: string[], type: BookType): Promise<GeneralResponse.ErrorItem | GeneralResponse.SuccessItem> {
try {
await this.InitService()
for (let index = 0; index < bookTaskDetailIds.length; index++) {
const bookTaskDetailId = bookTaskDetailIds[index];
let bookTaskDetail = this.bookTaskDetailService.GetBookTaskDetailDataById(bookTaskDetailId)
if (bookTaskDetail == null) {
throw new Error("没有找到对应的分镜数据")
}
let book = this.bookService.GetBookDataById(bookTaskDetail.bookId)
// 是不是又额外的类型,没有的话试用小说的类型
let task_type = undefined as BookBackTaskType
if (!type) {
type = book.type
}
switch (type) {
case BookType.MJ_REVERSE:
task_type = BookBackTaskType.MJ_REVERSE
break;
case BookType.SD_REVERSE:
task_type = BookBackTaskType.SD_REVERSE
break;
default:
throw new Error("暂不支持的推理类型")
}
let taskRes = await this.bookBackTaskListService.AddBookBackTask(book.id, task_type, TaskExecuteType.AUTO, bookTaskDetail.bookTaskId, bookTaskDetail.id
);
if (taskRes.code == 0) {
throw new Error(taskRes.message)
}
// 添加返回日志
await this.taskScheduler.AddLogToDB(book.id, book.type, `添加 ${task_type} 反推任务成功`, bookTaskDetail.bookTaskId, LoggerStatus.SUCCESS)
}
return successMessage(null, "添加反推任务成功", "ReverseBook_AddReversePromptTask")
} catch (error) {
return errorMessage("添加单个反推失败,错误信息如下:" + error.message, "ReverseBook_SingleReversePrompt")
}
}
// TODO SD反推待重写
async SDReversePrompt(task: TaskModal.Task): Promise<GeneralResponse.ErrorItem | GeneralResponse.SuccessItem> {
return successMessage(null, "", "ReverseBook_SDReversePrompt")
}
/**
* 执行单句推理的任务
* @param task 任务处理
*/
async SingleReversePrompt(task: TaskModal.Task): Promise<GeneralResponse.ErrorItem | GeneralResponse.SuccessItem> {
try {
// 开始处理
let res = undefined
switch (task.type) {
case BookBackTaskType.MJ_REVERSE:
res = await this.mjOpt.MJImage2Text(task);
break
case BookBackTaskType.SD_REVERSE:
res = await this.SDReversePrompt(task);
break
default:
throw new Error("未知的任务类型")
}
// 正常返回信息处理
return successMessage(null, `${task.name} _ ${task.id} 反推任务完成`, 'ReverseBook_SingleReversePrompt')
} catch (error) {
// 进行错误信息处理
return errorMessage("单个反推失败,错误信息如下:" + error.message, "ReverseBook_SingleReversePrompt")
}
}
/**
* 删除指定的提示词数据
* @param bookTaskDetailIds 要删除的提示词ID
*/
async RemoveReverseData(bookTaskDetailIds: string[]) {
try {
await this.InitService()
// 开始删除
if (bookTaskDetailIds.length <= 0) {
throw new Error("没有传入要删除的数据")
}
for (let i = 0; i < bookTaskDetailIds.length; i++) {
const element = bookTaskDetailIds[i];
this.bookTaskDetailService.DeleteBookTaskDetailReversePromptById(element)
}
// 全部删除完毕
return successMessage(null, "删除反推数据成功", "ReverseBook_RemoveReverseData")
} catch (error) {
return errorMessage("删除反推数据失败,错误信息如下:" + error.message, "ReverseBook_RemoveReverseData")
}
}
//#endregion
}