From 9036d5cf13e64552e8bbb79647a5619327dce99f Mon Sep 17 00:00:00 2001 From: Julian Waller Date: Tue, 10 Sep 2024 12:34:04 +0100 Subject: [PATCH] feat: downloadStill from media pool (#167) --- src/atem.ts | 45 +++++- .../DataTransferDownloadRequestCommand.ts | 2 +- src/dataTransfer/dataTransferDownloadStill.ts | 84 +++++++++++ src/dataTransfer/index.ts | 7 + src/lib/converters/__tests__/rle.spec.ts | 137 +++++++++++++++++- src/lib/converters/colorConstants.ts | 6 + src/lib/converters/rle.ts | 31 ++++ src/lib/converters/yuv422ToRgba.ts | 64 ++++++++ 8 files changed, 373 insertions(+), 3 deletions(-) create mode 100644 src/dataTransfer/dataTransferDownloadStill.ts create mode 100644 src/lib/converters/yuv422ToRgba.ts diff --git a/src/atem.ts b/src/atem.ts index 99581b927..2f9c7c6d6 100644 --- a/src/atem.ts +++ b/src/atem.ts @@ -19,7 +19,7 @@ import { InputChannel } from './state/input' import { DownstreamKeyerGeneral, DownstreamKeyerMask } from './state/video/downstreamKeyers' import * as DT from './dataTransfer' import * as Util from './lib/atemUtil' -import { getVideoModeInfo } from './lib/videoMode' +import { VideoModeInfo, getVideoModeInfo } from './lib/videoMode' import * as Enums from './enums' import { ClassicAudioMonitorChannel, @@ -56,6 +56,8 @@ import { TimeInfo } from './state/info' import { SomeAtemAudioLevels } from './state/levels' import { generateUploadBufferInfo, UploadBufferInfo } from './dataTransfer/dataTransferUploadBuffer' import { convertWAVToRaw } from './lib/converters/wavAudio' +import { decodeRLE } from './lib/converters/rle' +import { convertYUV422ToRGBA } from './lib/converters/yuv422ToRgba' export interface AtemOptions { address?: string @@ -147,6 +149,15 @@ export class BasicAtem extends EventEmitter { return this._state } + /** + * Get the current videomode of the ATEM, if known + */ + get videoMode(): Readonly | undefined { + if (!this.state) return undefined + + return getVideoModeInfo(this.state.settings.videoMode) + } + public async connect(address: string, port?: number): Promise { return this.socket.connect(address, port) } @@ -757,6 +768,38 @@ export class Atem extends BasicAtem { return this.sendCommand(command) } + /** + * Download a still image from the ATEM media pool + * + * Note: This performs colour conversions in JS, which is not very CPU efficient. If performance is important, + * consider using [@atem-connection/image-tools](https://www.npmjs.com/package/@atem-connection/image-tools) to + * pre-convert the images with more optimal algorithms + * @param index Still index to download + * @param format The pixel format to return for the downloaded image. 'raw' passes through unchanged, and will be RLE encoded. + * @returns Promise which returns the image once downloaded. If the still slot is not in use, this will throw + */ + public async downloadStill(index: number, format: 'raw' | 'rgba' | 'yuv' = 'rgba'): Promise { + let rawBuffer = await this.dataTransferManager.downloadStill(index) + + if (format === 'raw') { + return rawBuffer + } + + if (!this.state) throw new Error('Unable to check current resolution') + const resolution = getVideoModeInfo(this.state.settings.videoMode) + if (!resolution) throw new Error('Failed to determine required resolution') + + rawBuffer = decodeRLE(rawBuffer, resolution.width * resolution.height * 4) + + switch (format) { + case 'yuv': + return rawBuffer + case 'rgba': + default: + return convertYUV422ToRGBA(resolution.width, resolution.height, rawBuffer) + } + } + /** * Upload a still image to the ATEM media pool * diff --git a/src/commands/DataTransfer/DataTransferDownloadRequestCommand.ts b/src/commands/DataTransfer/DataTransferDownloadRequestCommand.ts index cdee6036e..8750bcfb0 100644 --- a/src/commands/DataTransfer/DataTransferDownloadRequestCommand.ts +++ b/src/commands/DataTransfer/DataTransferDownloadRequestCommand.ts @@ -16,7 +16,7 @@ export class DataTransferDownloadRequestCommand extends BasicWritableCommand { + #data: Buffer[] = [] + + constructor(public readonly stillIndex: number) { + super() + } + + public async startTransfer(transferId: number): Promise { + const command = new DataTransferDownloadRequestCommand({ + transferId: transferId, + transferStoreId: 0x00, + transferIndex: this.stillIndex, + transferType: 0x00f9, + }) + + return { + newState: DataTransferState.Ready, + commands: [command], + } + } + + public async handleCommand( + command: IDeserializedCommand, + oldState: DataTransferState + ): Promise { + if (command instanceof DataTransferErrorCommand) { + switch (command.properties.errorCode) { + case ErrorCode.Retry: + return this.restartTransfer(command.properties.transferId) + + case ErrorCode.NotFound: + this.abort(new Error('Invalid download')) + + return { + newState: DataTransferState.Finished, + commands: [], + } + default: + // Abort the transfer. + this.abort(new Error(`Unknown error ${command.properties.errorCode}`)) + + return { + newState: DataTransferState.Finished, + commands: [], + } + } + } else if (command instanceof DataTransferDataCommand) { + this.#data.push(command.properties.body) + + // todo - have we received all data? maybe check if the command.body < max_len + + return { + newState: oldState, + commands: [ + new DataTransferAckCommand({ + transferId: command.properties.transferId, + transferIndex: this.stillIndex, + }), + ], + } + } else if (command instanceof DataTransferCompleteCommand) { + this.resolvePromise(Buffer.concat(this.#data)) + + return { + newState: DataTransferState.Finished, + commands: [], + } + } + + return { newState: oldState, commands: [] } + } +} diff --git a/src/dataTransfer/index.ts b/src/dataTransfer/index.ts index 8e8d72bc4..f683dba14 100644 --- a/src/dataTransfer/index.ts +++ b/src/dataTransfer/index.ts @@ -11,6 +11,7 @@ import { DataTransferUploadMacro } from './dataTransferUploadMacro' import { LockObtainedCommand, LockStateUpdateCommand } from '../commands/DataTransfer' import debug0 from 'debug' import type { UploadBufferInfo } from './dataTransferUploadBuffer' +import { DataTransferDownloadStill } from './dataTransferDownloadStill' const MAX_PACKETS_TO_SEND_PER_TICK = 50 const MAX_TRANSFER_INDEX = (1 << 16) - 1 // Inclusive maximum @@ -170,6 +171,12 @@ export class DataTransferManager { } } + public async downloadStill(index: number): Promise { + const transfer = new DataTransferDownloadStill(index) + + return this.#stillsLock.enqueue(transfer) + } + public async uploadStill(index: number, data: UploadBufferInfo, name: string, description: string): Promise { const transfer = new DataTransferUploadStill(index, data, name, description) return this.#stillsLock.enqueue(transfer) diff --git a/src/lib/converters/__tests__/rle.spec.ts b/src/lib/converters/__tests__/rle.spec.ts index dbb819ec3..34d85a10c 100644 --- a/src/lib/converters/__tests__/rle.spec.ts +++ b/src/lib/converters/__tests__/rle.spec.ts @@ -1,4 +1,4 @@ -import { encodeRLE } from '../rle' +import { decodeRLE, encodeRLE } from '../rle' describe('encodeRLE', () => { test('no repetitions', () => { @@ -134,3 +134,138 @@ abababababababab` expect(encoded.toString('hex')).toEqual(expectation) }) }) + +describe('decodeRLE', () => { + test('no repetitions', () => { + const source = `abababababababab\ +cdcdcdcdcdcdcdcd\ +abababababababab\ +cdcdcdcdcdcdcdcd` + const decoded = decodeRLE(Buffer.from(source, 'hex'), source.length / 2) + expect(decoded.toString('hex')).toEqual(source) + }) + + test('two repetitions', () => { + const source = `abababababababab\ +abababababababab\ +cdcdcdcdcdcdcdcd\ +0000000000000000\ +1111111111111111` + const decoded = decodeRLE(Buffer.from(source, 'hex'), source.length / 2) + expect(decoded.toString('hex')).toEqual(source) + }) + + test('three repetitions', () => { + const source = `abababababababab\ +abababababababab\ +abababababababab\ +cdcdcdcdcdcdcdcd\ +0000000000000000\ +1111111111111111` + const decoded = decodeRLE(Buffer.from(source, 'hex'), source.length / 2) + expect(decoded.toString('hex')).toEqual(source) + }) + + test('four repetitions', () => { + const source = `fefefefefefefefe\ +0000000000000004\ +abababababababab\ +cdcdcdcdcdcdcdcd\ +0000000000000000\ +1111111111111111` + const expectation = `abababababababab\ +abababababababab\ +abababababababab\ +abababababababab\ +cdcdcdcdcdcdcdcd\ +0000000000000000\ +1111111111111111` + const decoded = decodeRLE(Buffer.from(source, 'hex'), expectation.length / 2) + expect(decoded.toString('hex')).toEqual(expectation) + }) + + test('five repetitions at the beginning', () => { + const source = `fefefefefefefefe\ +0000000000000005\ +abababababababab\ +cdcdcdcdcdcdcdcd\ +0000000000000000\ +1111111111111111` + const expectation = `abababababababab\ +abababababababab\ +abababababababab\ +abababababababab\ +abababababababab\ +cdcdcdcdcdcdcdcd\ +0000000000000000\ +1111111111111111` + const decoded = decodeRLE(Buffer.from(source, 'hex'), expectation.length / 2) + expect(decoded.toString('hex')).toEqual(expectation) + }) + + test('five repetitions in the midddle', () => { + const source = `2323232323232323\ +fefefefefefefefe\ +0000000000000005\ +abababababababab\ +cdcdcdcdcdcdcdcd\ +0000000000000000\ +1111111111111111` + const expectation = `2323232323232323\ +abababababababab\ +abababababababab\ +abababababababab\ +abababababababab\ +abababababababab\ +cdcdcdcdcdcdcdcd\ +0000000000000000\ +1111111111111111` + const decoded = decodeRLE(Buffer.from(source, 'hex'), expectation.length / 2) + expect(decoded.toString('hex')).toEqual(expectation) + }) + + test('five repetitions in the midddle #2', () => { + const source = `2323232323232323\ +fefefefefefefefe\ +0000000000000005\ +abababababababab\ +cdcdcdcdcdcdcdcd` + const expectation = `2323232323232323\ +abababababababab\ +abababababababab\ +abababababababab\ +abababababababab\ +abababababababab\ +cdcdcdcdcdcdcdcd` + const decoded = decodeRLE(Buffer.from(source, 'hex'), expectation.length / 2) + expect(decoded.toString('hex')).toEqual(expectation) + }) + + test('five repetitions at the end', () => { + const source = `2323232323232323\ +fefefefefefefefe\ +0000000000000005\ +abababababababab` + const expectation = `2323232323232323\ +abababababababab\ +abababababababab\ +abababababababab\ +abababababababab\ +abababababababab` + const decoded = decodeRLE(Buffer.from(source, 'hex'), expectation.length / 2) + expect(decoded.toString('hex')).toEqual(expectation) + }) + + test('only five repetitions', () => { + const source = `fefefefefefefefe\ +0000000000000005\ +abababababababab` + const expectation = `abababababababab\ +abababababababab\ +abababababababab\ +abababababababab\ +abababababababab` + const decoded = decodeRLE(Buffer.from(source, 'hex'), expectation.length / 2) + expect(decoded.toString('hex')).toEqual(expectation) + }) +}) diff --git a/src/lib/converters/colorConstants.ts b/src/lib/converters/colorConstants.ts index 67990ed3c..25d158c75 100644 --- a/src/lib/converters/colorConstants.ts +++ b/src/lib/converters/colorConstants.ts @@ -13,6 +13,9 @@ export interface ColorConvertConstants { readonly YOffset: number readonly CbCrOffset: number + readonly KRKRioKG: number + readonly KBKBioKG: number + readonly KRoKBi: number readonly KGoKBi: number readonly KBoKRi: number @@ -44,6 +47,9 @@ function createColorConvertConstants(KR: number, KB: number): ColorConvertConsta YOffset: 16 << 8, CbCrOffset: 128 << 8, + KRKRioKG: (KR * KRi * 2) / KG, + KBKBioKG: (KB * KBi * 2) / KG, + KRoKBi: (KR / KBi) * HalfCbCrRange, KGoKBi: (KG / KBi) * HalfCbCrRange, KBoKRi: (KB / KRi) * HalfCbCrRange, diff --git a/src/lib/converters/rle.ts b/src/lib/converters/rle.ts index f87f51092..927f73e1a 100644 --- a/src/lib/converters/rle.ts +++ b/src/lib/converters/rle.ts @@ -51,3 +51,34 @@ export function encodeRLE(data: Buffer): Buffer { return result.slice(0, resultOffset + 8) } + +export function decodeRLE(data: Buffer, fullSize: number): Buffer { + const result = Buffer.alloc(fullSize) + + let resultOffset = -8 + + for (let sourceOffset = 0; sourceOffset < data.length; sourceOffset += 8) { + const block = data.readBigUInt64BE(sourceOffset) + + // read a header, start a repeating block + if (block === RLE_HEADER) { + // Read the count + sourceOffset += 8 + const repeatCount = Number(data.readBigUInt64BE(sourceOffset)) + + // Read the repeated sample + sourceOffset += 8 + const repeatBlock = data.readBigUInt64BE(sourceOffset) + + // Write to the output + for (let i = 0; i < repeatCount; i++) { + result.writeBigUInt64BE(repeatBlock, (resultOffset += 8)) + } + } else { + // No RLE, repeat unchanged + result.writeBigUInt64BE(block, (resultOffset += 8)) + } + } + + return result +} diff --git a/src/lib/converters/yuv422ToRgba.ts b/src/lib/converters/yuv422ToRgba.ts new file mode 100644 index 000000000..e66fd56d6 --- /dev/null +++ b/src/lib/converters/yuv422ToRgba.ts @@ -0,0 +1,64 @@ +/** + * @todo: MINT - 2018-5-24: + * Create util functions that handle proper colour spaces in UHD. + */ + +import { ColorConvertConstantsBT709, ColorConvertConstantsBT601 } from './colorConstants' + +function clamp(v: number) { + if (v <= 0) return 0 + if (v >= 255) return 255 + return v +} + +export function convertYUV422ToRGBA(width: number, height: number, data: Buffer): Buffer { + const constants = height >= 720 ? ColorConvertConstantsBT709 : ColorConvertConstantsBT601 + + const splitSample = (raw: number): [y: number, uv: number, a: number] => { + const y = (raw & 0x000003ff) << 6 + const uv = (raw & 0x000ffc00) >> 4 // same as << 6 >> 10 + const a = (raw & 0x3ff00000) >> 20 + + const y_full = (y - constants.YOffset) / constants.YRange + const uv_full = (uv - constants.CbCrOffset) / constants.HalfCbCrRange + const a_full = (((a - (16 << 2)) / 219) * 255) >> 2 // TODO - confirm correct range + + return [y_full, uv_full, a_full] + } + + const genColor = ( + y8: number, + cb8: number, + cr8: number, + a10: number + ): [r: number, g: number, b: number, a: number] => { + const r = clamp(Math.round(y8 + constants.KRi * cr8)) + const g = clamp(Math.round(y8 - constants.KRKRioKG * cr8 - constants.KBKBioKG * cb8)) + const b = clamp(Math.round(y8 + constants.KBi * cb8)) + const a = Math.round(a10) + + return [r, g, b, a] + } + + const buffer = Buffer.alloc(width * height * 4) + for (let i = 0; i < width * height * 4; i += 8) { + const sample1 = data.readUint32BE(i) + const sample2 = data.readUint32BE(i + 4) + + const [y8a, cb8, a10a] = splitSample(sample1) + const [y8b, cr8, a10b] = splitSample(sample2) + + const [r1, g1, b1, a1] = genColor(y8a, cb8, cr8, a10a) + const [r2, g2, b2, a2] = genColor(y8b, cb8, cr8, a10b) + + buffer.writeUint8(r1, i) + buffer.writeUint8(g1, i + 1) + buffer.writeUint8(b1, i + 2) + buffer.writeUint8(a1, i + 3) + buffer.writeUint8(r2, i + 4) + buffer.writeUint8(g2, i + 5) + buffer.writeUint8(b2, i + 6) + buffer.writeUint8(a2, i + 7) + } + return buffer +}