LaiTool/src/main/Service/watermark.ts

import path from 'path'
import fs from 'fs'
import util from 'util'
import { exec } from 'child_process'
const execAsync = util.promisify(exec);
import { CheckFileOrDirExist, CheckFolderExistsOrCreate, CopyFileOrFolder, GetFilesWithExtensions } from '../../define/Tools/file'
import { errorMessage, successMessage } from '../Public/generalTools'
import { SoftwareService } from '../../define/db/service/SoftWare/softwareService'
import { isEmpty } from 'lodash'
import { RemoveWatermarkType } from '../../define/enum/waterMarkAndSubtitle'
import { ValidateJson } from '../../define/Tools/validate'
import { define } from '../../define/define'
import { LOGGER_DEFINE } from '../../define/logger_define'
import axios from 'axios'
import { Base64ToFile, GetImageBase64 } from '../../define/Tools/image'
import { TaskScheduler } from './taskScheduler';
import { LoggerStatus, OtherData, ResponseMessageType } from '../../define/enum/softwareEnum';
import { basicApi } from '../../api/apiBasic';
import { FfmpegOptions } from './ffmpegOptions';
import { ProcessImage } from '../../define/Tools/image';
import { BookService } from '../../define/db/service/Book/bookService';
import { OperateBookType } from '../../define/enum/bookEnum';
import { GeneralResponse } from '../../model/generalResponse';
import { BookTaskDetailService } from '../../define/db/service/Book/bookTaskDetailService';
import { Book } from '../../model/book';
import { DEFINE_STRING } from '../../define/define_string';
import { BookTaskService } from '../../define/db/service/Book/bookTaskService';

export class Watermark {
  softwareService: SoftwareService
  taskScheduler: TaskScheduler;
  bookService: BookService
  bookTaskDetailService: BookTaskDetailService
  bookTaskService: BookTaskService
  constructor() {
  }

  async InitService() {
    if (!this.softwareService) {
      this.softwareService = await SoftwareService.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.bookTaskService) {
      this.bookTaskService = await BookTaskService.getInstance()
    }
  }

  // 主动返回前端的消息
  sendReturnMessage(data: GeneralResponse.MessageResponse, message_name = DEFINE_STRING.BOOK.GET_COPYWRITING_RETURN) {
    global.newWindow[0].win.webContents.send(message_name, data)
  }

  //#region  设置

  /**
   * 初始化水印设置
   */
  private InitWatermarkSetting() {
    let defaultData = {
      selectModel: RemoveWatermarkType.LOCAL_LAMA,
      iopaint: {
        url: "http://127.0.0.1:8080/"
      }
    } as ImageModel.RemoveWatermarkSetting

    this.softwareService.SaveSoftwarePropertyData('watermarkSetting', JSON.stringify(defaultData))
    return defaultData;
  }

  /**
   * 获取去除水印的设置
   */
  async GetWatermarkSetting() {
    try {
      await this.InitService();
      let watermarkSetting = undefined;
      let setting = this.softwareService.GetSoftWarePropertyData("watermarkSetting");
      if (!isEmpty(setting)) {
        // 初始化
        let tryP = ValidateJson(setting);
        if (!tryP) {
          throw new Error('解析去除水印的设置失败，数据格式不正确')
        }
        watermarkSetting = JSON.parse(setting);
      } else {
        watermarkSetting = this.InitWatermarkSetting()
      }
      return successMessage(watermarkSetting, '获取去除水印的设置成功', 'Image_GetWatermarkSetting')
    } catch (error) {
      return errorMessage("获取去除水印的设置失败，失败信息如下：" + error.toString(), 'Image_GetWatermarkSetting')
    }
  }

  /**
   * 保存去除水印的设置
   * @param value
   * @returns 
   */
  async SaveWatermarkSetting(value: ImageModel.RemoveWatermarkSetting) {
    try {
      await this.InitService();
      if (value.selectModel == RemoveWatermarkType.IOPAINT && !value.iopaint.url) {
        throw new Error('iopaint模式，iopaint的地址不能为空')
      }

      if (value.selectModel == RemoveWatermarkType.LOCAL_LAMA) {
        // 简单判断环境是不是存在
        let lamaPath = path.join(define.scripts_path, 'lama/lama_inpaint.exe');
        if (!(await CheckFileOrDirExist(lamaPath))) {
          throw new Error('本地LAMA环墋不存在，请先安装，详细信息看下方的 ‘lama安装教程’')
        }
      }

      if (value.iopaint.url) {
        // 验证iopaint的地址是不是一个合法的url
        let reg = new RegExp('^(http|https)://')
        if (!reg.test(value.iopaint.url)) {
          throw new Error('iopaint的地址不是一个合法的网络地址')
        }
        if (!value.iopaint.url.endsWith('/')) {
          value.iopaint.url = value.iopaint.url + '/'
        }
      }


      // 保存数据
      let res = this.softwareService.SaveSoftwarePropertyData('watermarkSetting', JSON.stringify(value));
      return successMessage(null, "保存去除水印的设置成功", 'Image_SaveWatermarkSetting')
    } catch (error) {
      return errorMessage("保存去除水印的设置失败，失败信息如下：" + error.toString(), 'Image_SaveWatermarkSetting')
    }
  }

  //#endregion

  //#region 去除水印相关操作

  /**
   * 去除水印，试用本地lama
   * @param value 
   */
  async ProcessImageByLama(value: ImageModel.ProcessImageParams): Promise<Buffer | string> {
    try {
      let lama_script = path.resolve(define.scripts_path, `lama/lama_inpaint.exe`)
      // 就是判断指定的文件和文件夹是不是存在
      let has_exe = await CheckFileOrDirExist(lama_script)
      if (!has_exe) {
        throw new Error('图片水印处理组件不存在，请看教程自行下载')
      }
      let has_model = await CheckFileOrDirExist(
        path.resolve(define.scripts_path, 'lama/model/big-lama.pt')
      )
      if (!has_model) {
        throw new Error('图片水印处理的模型不存在，请看教程自行下载')
      }
      global.logger.info(
        LOGGER_DEFINE.REMOVE_WATERMARK,
        `开始使用lama去除水印，开始调用lama程序`
      )
      let script = `cd "${path.dirname(lama_script)}" && "${lama_script}" "-l" "${value.inputFilePath}" "${value.maskPath}" "${value.outFilePath}"`
      let scriptRes = await execAsync(script, { maxBuffer: 1024 * 1024 * 10, encoding: 'utf-8' })
      // 判断是不是有输出文件
      let has_out = await CheckFileOrDirExist(value.outFilePath)
      if (!has_out) {
        throw new Error('lama去除水印失败，没有输出文件')
      }

      if (value.type == 'arrayBuffer') {
        // 读取导出的文件
        let res_data = await fs.promises.readFile(value.outFilePath)
        return res_data
      } else if (value.type == 'file') {
        return value.outFilePath
      }

    } catch (error) {
      throw new Error('使用Lama去除水印失败，失败信息如下：' + error.toString())
    }
  }

  /**
   * 去除s水印，使用iopaint 
   * @param value 传入的参数
   * @param iopaintUrl 
   */
  async ProcessImageByIopaint(value: ImageModel.ProcessImageParams, iopaintUrl: string): Promise<Buffer | string> {
    try {
      iopaintUrl = iopaintUrl + 'api/v1/inpaint'
      let headers = {
        accept: '*/*',
        'accept-language': 'zh-CN,zh;q=0.9',
        'content-type': 'application/json'
      }
      let data = {
        image: value.imageBase64,
        mask: value.maskBase64,
        ldm_steps: 30,
        ldm_sampler: 'ddim',
        zits_wireframe: true,
        cv2_flag: 'INPAINT_NS',
        cv2_radius: 5,
        hd_strategy: 'Crop',
        hd_strategy_crop_triger_size: 640,
        hd_strategy_crop_margin: 128,
        hd_trategy_resize_imit: 2048 * 5,
        prompt: '',
        negative_prompt:
          'out of frame, lowres, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, disfigured, gross proportions, malformed limbs, watermark, signature',
        use_croper: false,
        croper_x: 284,
        croper_y: 284,
        croper_height: 512,
        croper_width: 512,
        use_extender: false,
        extender_x: 0,
        extender_y: 0,
        extender_height: 1080,
        extender_width: 1080,
        sd_mask_blur: 12,
        sd_strength: 1,
        sd_steps: 50,
        sd_guidance_scale: 7.5,
        sd_sampler: 'DPM++ 2M',
        sd_seed: -1,
        sd_match_histograms: false,
        sd_lcm_lora: false,
        paint_by_example_example_image: null,
        p2p_image_guidance_scale: 1.5,
        enable_controlnet: false,
        controlnet_conditioning_scale: 0.4,
        controlnet_method: '',
        enable_brushnet: false,
        brushnet_method: 'random_mask',
        brushnet_conditioning_scale: 1,
        enable_powerpaint_v2: false,
        powerpaint_task: 'text-guided'
      }
      let res = await basicApi.post(iopaintUrl, data, headers)

      if (value.type == 'arrayBuffer') {
        return res.data
      } else if (value.type == 'file') {
        let buffer = Buffer.from(res.data)
        await fs.promises.writeFile(value.outFilePath, buffer)
        return value.outFilePath
      }
    } catch (error) {
      throw new Error('使用Iopaint去除水印失败，失败信息如下：' + error.toString())
    }
  }

  /**
   * 去除图片水印，当前只支持本地，后续会支持 iopaint 
   * @param value 
   */
  async ProcessImage(value: ImageModel.ProcessImageParams): Promise<string | Buffer> {
    let outDir = path.dirname(value.outFilePath)
    await CheckFolderExistsOrCreate(outDir);

    let watermarkSettingRes = await this.GetWatermarkSetting();
    if (watermarkSettingRes.code == 0) {
      throw new Error(watermarkSettingRes.message)
    }
    let watermarkSetting = watermarkSettingRes.data as ImageModel.RemoveWatermarkSetting;

    switch (watermarkSetting.selectModel) {
      case RemoveWatermarkType.LOCAL_LAMA:
        return await this.ProcessImageByLama(value)
      case RemoveWatermarkType.IOPAINT:
        return await this.ProcessImageByIopaint(value, watermarkSetting.iopaint.url)
      default:
        throw new Error('未知的去除水印模式')
    }
  }
  //#endregion

  //#region 开始去除水印


  /**
 * 查看去除水印设置 去除图片指定位置的水印，这个是要单独做蒙板处理的
 * @param imageBase64 基础图片的base64
 * @param maskPosition 前端勾选的水印的位置
 */
  async ProcessImageCheck(imageBase64: string, maskPosition: [], bookId: string) {
    try {
      await this.InitService()
      // 开始做处理
      console.log(imageBase64, maskPosition);
      let book = this.bookService.GetBookDataById(bookId)
      if (book == null) {
        throw new Error('指定的小说数据不存在')
      }
      let inputFilePath = path.resolve(book.bookFolderPath, `data/mask/temp/${new Date().getTime()}.png`)
      await CheckFolderExistsOrCreate(path.dirname(inputFilePath))
      const base64Image = imageBase64.split(';base64,').pop();
      await fs.promises.writeFile(inputFilePath, base64Image, { encoding: 'base64' })
      let outputPath = path.resolve(book.bookFolderPath, `data/mask/mask_temp_${new Date().getTime()}.png`)
      let outImagePath = path.resolve(book.bookFolderPath, `data/mask/temp/out_${new Date().getTime()}.png`)

      await ProcessImage(inputFilePath, outputPath, maskPosition.map((item: any) => { return { x: item.startX, y: item.startY, width: item.width, height: item.height } }))
      let markBase64 = await GetImageBase64(outputPath)

      // 开始去谁赢
      let res = await this.ProcessImage({
        imageBase64: imageBase64,
        maskBase64: markBase64, // 处理蒙板的base64
        type: 'arrayBuffer', // 返回数据的类型（是直接写道输出文件地址，还是返回Buffer数组）
        inputFilePath: inputFilePath, // 输入文件的地址（待处理的）
        maskPath: outputPath, // 蒙板文件的地址
        outFilePath: outImagePath // 输出文件的地址
      })

      // 将得到的buffer返回前端进行渲染
      return successMessage(res, "去除水印成功", "Image_ProcessImageCheck")

    } catch (error) {
      return errorMessage('去除图片指定位置的水印失败，失败信息如下：' + error.toString(), 'Image_ProcessImageCheck')
    }

  }

  /**
   * 执行去除水印操作
   * @param bookId 小说ID
   */
  async RemoveWatermark(id: string, operateBookType: OperateBookType): Promise<GeneralResponse.ErrorItem | GeneralResponse.SuccessItem> {
    try {
      await this.InitService()
      let book = undefined as Book.SelectBook
      let bookTask = undefined as Book.SelectBookTask
      let bookTaskDetails = undefined as Book.SelectBookTaskDetail[]
      if (operateBookType == OperateBookType.BOOKTASK) {
        bookTask = this.bookTaskService.GetBookTaskDataById(id);
        if (bookTask == null) {
          throw new Error('指定的小说任务数据不存在')
        }
        book = this.bookService.GetBookDataById(bookTask.bookId)
        if (book == null) {
          throw new Error("小说数据不存在，请检查")
        }
        bookTaskDetails = this.bookTaskDetailService.GetBookTaskData({
          bookTaskId: bookTask.id,
          bookId: bookTask.bookId
        }).data as Book.SelectBookTaskDetail[]
      } else if (operateBookType == OperateBookType.BOOKTASKDETAIL) {
        let bookTaskDetail = this.bookTaskDetailService.GetBookTaskDetailDataById(id)
        if (bookTaskDetail == null) {
          throw new Error("指定的小说任务分镜信息不存在，请检查")
        }
        bookTask = this.bookTaskService.GetBookTaskDataById(bookTaskDetail.bookTaskId);
        if (bookTask == null) {
          throw new Error('指定的小说任务数据不存在')
        }
        book = this.bookService.GetBookDataById(bookTask.bookId)
        if (book == null) {
          throw new Error("小说数据不存在，请检查")
        }

        bookTaskDetails = [bookTaskDetail];
      } else {
        throw new Error("未知的操作类型")
      }

      let watermarkPosition = book.watermarkPosition;
      if (isEmpty(watermarkPosition)) {
        throw new Error("当前没有文案位置的内容，请先进行 ‘开始去除水印 -> 选择水印位置 -> 框选大概位置 -> 保存蒙板位置’ 操作保存蒙板的位置")
      }
      if (!ValidateJson(watermarkPosition)) {
        throw new Error("文案位置格式有误，请先进行 ‘开始去除水印 -> 选择水印位置 -> 框选大概位置 -> 保存蒙板位置’ 操作保存蒙板的位置")
      }
      let watermarkPositionParse = JSON.parse(watermarkPosition) as any[]
      if (watermarkPositionParse.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));
      // 这边要计算倍率 
      let watermarkPositionArray = watermarkPositionParse.map(item => {
        return {
          x: item.startX,
          y: item.startY,
          width: item.width,
          height: item.height,
          imageWidth: item.imageWidth,
          imageHeight: item.imageHeight
        }
      })
      await ProcessImage(inputImage, outImagePath, watermarkPositionArray)

      if (!(await CheckFileOrDirExist(outImagePath))) {
        throw new Error("生成蒙板失败")
      }

      // 开始去除水印
      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)
        // 新的输出地址
        let tempOutImagePath = path.join(path.dirname(bookInputImage), "temp_" + path.basename(outImagePath))
        // await fs.promises.unlink(bookInputImage) // 删除原来的图片
        // 开始调用去除水印的方法
        let imageBase64 = await GetImageBase64(bak)
        let maskBase64 = await GetImageBase64(outImagePath)

        let res = await this.ProcessImage({
          imageBase64: imageBase64,
          maskBase64: maskBase64,
          type: 'file',
          inputFilePath: bak,
          maskPath: outImagePath,
          outFilePath: tempOutImagePath
        })
        // 执行完成，将缓存的图片删除，复制新的图片到原来的位置
        await fs.promises.unlink(bookInputImage)
        await CopyFileOrFolder(tempOutImagePath, bookInputImage)
        // 删除临时文件
        await fs.promises.unlink(tempOutImagePath)

        // 去水印执行完毕
        this.sendReturnMessage({
          code: 1,
          id: element.id,
          type: ResponseMessageType.REMOVE_WATERMARK,
          data: bookInputImage
        }, DEFINE_STRING.BOOK.REMOVE_WATERMARK_RETURN)


        this.taskScheduler.AddLogToDB(book.id, book.type, `${element.name} 去除水印完成`, element.bookTaskId, LoggerStatus.SUCCESS)
      }
      // 全部完毕
      if (operateBookType == OperateBookType.BOOKTASKDETAIL) {
        return successMessage(bookTaskDetails[0].oldImage + '?t=' + new Date().getTime(), "去除水印完成", "ReverseBook_RemoveWatermark")
      } else {
        return successMessage(null, "全部图片去除水印完成", "ReverseBook_RemoveWatermark")
      }
    } catch (error) {
      return errorMessage("去除水印失败，错误信息如下：" + error.message, "ReverseBook_RemoveWatermark")
    }
  }

  //#endregion
}