PeerTube/server/lib/job-queue/handlers/generate-storyboard.ts

import { Job } from 'bullmq'
import { join } from 'path'
import { getFFmpegCommandWrapperOptions } from '@server/helpers/ffmpeg'
import { generateImageFilename, getImageSize } from '@server/helpers/image-utils'
import { logger, loggerTagsFactory } from '@server/helpers/logger'
import { CONFIG } from '@server/initializers/config'
import { STORYBOARD } from '@server/initializers/constants'
import { federateVideoIfNeeded } from '@server/lib/activitypub/videos'
import { VideoPathManager } from '@server/lib/video-path-manager'
import { StoryboardModel } from '@server/models/video/storyboard'
import { VideoModel } from '@server/models/video/video'
import { MVideo } from '@server/types/models'
import { FFmpegImage, isAudioFile } from '@shared/ffmpeg'
import { GenerateStoryboardPayload } from '@shared/models'

const lTagsBase = loggerTagsFactory('storyboard')

async function processGenerateStoryboard (job: Job): Promise<void> {
  const payload = job.data as GenerateStoryboardPayload
  const lTags = lTagsBase(payload.videoUUID)

  logger.info('Processing generate storyboard of %s in job %s.', payload.videoUUID, job.id, lTags)

  const video = await VideoModel.loadFull(payload.videoUUID)
  if (!video) {
    logger.info('Video %s does not exist anymore, skipping storyboard generation.', payload.videoUUID, lTags)
    return
  }

  const inputFile = video.getMaxQualityFile()

  await VideoPathManager.Instance.makeAvailableVideoFile(inputFile, async videoPath => {
    const isAudio = await isAudioFile(videoPath)

    if (isAudio) {
      logger.info('Do not generate a storyboard of %s since the video does not have a video stream', payload.videoUUID, lTags)
      return
    }

    const ffmpeg = new FFmpegImage(getFFmpegCommandWrapperOptions('thumbnail'))

    const filename = generateImageFilename()
    const destination = join(CONFIG.STORAGE.STORYBOARDS_DIR, filename)

    const totalSprites = buildTotalSprites(video)
    if (totalSprites === 0) {
      logger.info('Do not generate a storyboard of %s because the video is not long enough', payload.videoUUID, lTags)
      return
    }

    const spriteDuration = Math.round(video.duration / totalSprites)

    const spritesCount = findGridSize({
      toFind: totalSprites,
      maxEdgeCount: STORYBOARD.SPRITES_MAX_EDGE_COUNT
    })

    logger.debug(
      'Generating storyboard from video of %s to %s', video.uuid, destination,
      { ...lTags, spritesCount, spriteDuration, videoDuration: video.duration }
    )

    await ffmpeg.generateStoryboardFromVideo({
      destination,
      path: videoPath,
      sprites: {
        size: STORYBOARD.SPRITE_SIZE,
        count: spritesCount,
        duration: spriteDuration
      }
    })

    const imageSize = await getImageSize(destination)

    const existing = await StoryboardModel.loadByVideo(video.id)
    if (existing) await existing.destroy()

    await StoryboardModel.create({
      filename,
      totalHeight: imageSize.height,
      totalWidth: imageSize.width,
      spriteHeight: STORYBOARD.SPRITE_SIZE.height,
      spriteWidth: STORYBOARD.SPRITE_SIZE.width,
      spriteDuration,
      videoId: video.id
    })

    logger.info('Storyboard generation %s ended for video %s.', destination, video.uuid, lTags)
  })

  if (payload.federate) {
    await federateVideoIfNeeded(video, false)
  }
}

// ---------------------------------------------------------------------------

export {
  processGenerateStoryboard
}

function buildTotalSprites (video: MVideo) {
  const maxSprites = STORYBOARD.SPRITE_SIZE.height * STORYBOARD.SPRITE_SIZE.width
  const totalSprites = Math.min(Math.ceil(video.duration), maxSprites)

  // We can generate a single line
  if (totalSprites <= STORYBOARD.SPRITES_MAX_EDGE_COUNT) return totalSprites

  return findGridFit(totalSprites, STORYBOARD.SPRITES_MAX_EDGE_COUNT)
}

function findGridSize (options: {
  toFind: number
  maxEdgeCount: number
}) {
  const { toFind, maxEdgeCount } = options

  for (let i = 1; i <= maxEdgeCount; i++) {
    for (let j = i; j <= maxEdgeCount; j++) {
      if (toFind === i * j) return { width: j, height: i }
    }
  }

  throw new Error(`Could not find grid size (to find: ${toFind}, max edge count: ${maxEdgeCount}`)
}

function findGridFit (value: number, maxMultiplier: number) {
  for (let i = value; i--; i > 0) {
    if (!isPrimeWithin(i, maxMultiplier)) return i
  }

  throw new Error('Could not find prime number below ' + value)
}

function isPrimeWithin (value: number, maxMultiplier: number) {
  if (value < 2) return false

  for (let i = 2, end = Math.min(Math.sqrt(value), maxMultiplier); i <= end; i++) {
    if (value % i === 0 && value / i <= maxMultiplier) return false
  }

  return true
}