labelmain-demo/app/component/label/video2image/processVideo.ts

"use client"

import type { WorkerMessage, WorkerResponse } from "../ffmpeg/ffmpeg.worker"

export interface VideoFrameFile {
  name: string
  dataUrl: string
}

interface ProcessVideoParams {
  base64Data: string
  fileName: string
  namePrefix?: string
  baseURL?: string
  args?: string[]
}

const defaultArgs = [
  "-vf",
  "fps=1,scale='min(720,iw)':-1",
  "-q:v",
  "2",
  "frame_%03d.jpg",
]

const mimeByExt: Record<string, string> = {
  jpg: "image/jpeg",
  jpeg: "image/jpeg",
  png: "image/png",
  webp: "image/webp",
  mp4: "video/mp4",
  mov: "video/quicktime",
  webm: "video/webm",
  mkv: "video/x-matroska",
  avi: "video/x-msvideo",
  h264: "video/h264",
  "264": "video/h264",
}

const getFileExt = (fileName: string) => {
  const segments = fileName.split(".")
  if (segments.length <= 1) return ""
  return segments[segments.length - 1].toLowerCase()
}

const stripFileExt = (fileName: string) => {
  const index = fileName.lastIndexOf(".")
  if (index <= 0) return fileName
  return fileName.slice(0, index)
}

const normalizeBase64 = (text: string) => {
  if (!text) return ""
  const raw = text.includes(",") ? text.split(",").slice(1).join(",") : text
  return raw.replace(/\s+/g, "")
}

const sanitizePrefix = (prefix: string) => {
  const safe = prefix
    .replace(/[^a-zA-Z0-9_-]+/g, "_")
    .replace(/_+/g, "_")
    .replace(/^_|_$/g, "")
  return safe || "video"
}

const uint8ArrayToBase64 = (arr: Uint8Array) => {
  const chunkSize = 0x8000
  let binary = ""
  for (let i = 0; i < arr.length; i += chunkSize) {
    binary += String.fromCharCode(...arr.subarray(i, i + chunkSize))
  }
  return btoa(binary)
}

const base64ToFile = (base64Data: string, fileName: string) => {
  const payload = normalizeBase64(base64Data)
  if (!payload) throw new Error("empty video payload")
  const binary = atob(payload)
  const len = binary.length
  const bytes = new Uint8Array(len)
  for (let i = 0; i < len; i += 1) {
    bytes[i] = binary.charCodeAt(i)
  }
  const fileExt = getFileExt(fileName)
  return new File([bytes], fileName, {
    type: mimeByExt[fileExt] || "video/mp4",
  })
}

const getOutputExtFromArgs = (args: string[]) => {
  for (let i = args.length - 1; i >= 0; i -= 1) {
    const ext = getFileExt(args[i])
    if (ext) return ext
  }
  return "jpg"
}

export const processVideo = async (
  params: ProcessVideoParams
): Promise<VideoFrameFile[]> => {
  if (typeof window === "undefined") return []

  const args = params.args && params.args.length ? params.args : defaultArgs
  const outputExt = getOutputExtFromArgs(args)
  const prefix = sanitizePrefix(
    params.namePrefix || stripFileExt(params.fileName)
  )
  const baseURL = params.baseURL || `${window.location.origin}/wasm`
  const inputFile = base64ToFile(params.base64Data, params.fileName)

  const worker = new Worker(
    new URL("../ffmpeg/ffmpeg.worker.ts", import.meta.url)
  )
  let frameIndex = 0
  const frames: VideoFrameFile[] = []

  try {
    return await new Promise<VideoFrameFile[]>((resolve, reject) => {
      let settled = false

      const finish = (next: () => void) => {
        if (settled) return
        settled = true
        next()
      }

      worker.onmessage = (event: MessageEvent<WorkerResponse>) => {
        const message = event.data
        switch (message.type) {
          case "READY": {
            const runMessage: WorkerMessage = {
              type: "EXEC",
              file: inputFile,
              args,
              outputPattern: "frame_",
            }
            worker.postMessage(runMessage)
            return
          }
          case "SEGMENT_DATA": {
            const files = message.files || []
            files.forEach((file) => {
              const ext = getFileExt(file.name) || outputExt || "jpg"
              const frameName = `${prefix}_frame_${String(frameIndex).padStart(
                6,
                "0"
              )}.${ext}`
              frameIndex += 1
              frames.push({
                name: frameName,
                dataUrl: `data:${mimeByExt[ext] || "image/jpeg"};base64,${uint8ArrayToBase64(file.data)}`,
              })
            })
            return
          }
          case "DONE": {
            finish(() => {
              if (!frames.length) {
                reject(new Error("video has no extracted frames"))
              } else {
                resolve(frames)
              }
            })
            return
          }
          case "FATAL": {
            finish(() =>
              reject(new Error(message.error || "video process fatal"))
            )
            return
          }
          case "ERROR": {
            finish(() =>
              reject(new Error(message.error || "video process error"))
            )
            return
          }
          default:
            return
        }
      }

      worker.onerror = (event) => {
        finish(() => reject(new Error(event.message || "video worker error")))
      }

      const loadMessage: WorkerMessage = {
        type: "LOAD",
        baseURL,
      }
      worker.postMessage(loadMessage)
    })
  } finally {
    worker.terminate()
  }
}