feat: downloadStill from media pool (#167)

Author: Julusian

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<AtemEvents> {
 		return this._state
 	}
 
+	/**
+	 * Get the current videomode of the ATEM, if known
+	 */
+	get videoMode(): Readonly<VideoModeInfo> | undefined {
+		if (!this.state) return undefined
+
+		return getVideoModeInfo(this.state.settings.videoMode)
+	}
+
 	public async connect(address: string, port?: number): Promise<void> {
 		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<Buffer> {
+		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
 	 *

src/commands/DataTransfer/DataTransferDownloadRequestCommand.ts

@@ -16,7 +16,7 @@ export class DataTransferDownloadRequestCommand extends BasicWritableCommand<Dat
 		buffer.writeUInt16BE(this.properties.transferStoreId, 2)
 		buffer.writeUInt16BE(this.properties.transferIndex, 6)
 
-		buffer.writeUInt8(this.properties.transferType, 8)
+		buffer.writeUInt16BE(this.properties.transferType, 8)
 
 		return buffer
 	}

src/dataTransfer/dataTransferDownloadStill.ts

@@ -0,0 +1,84 @@
+import {
+	DataTransferAckCommand,
+	DataTransferCompleteCommand,
+	DataTransferDataCommand,
+	DataTransferDownloadRequestCommand,
+	DataTransferErrorCommand,
+	ErrorCode,
+} from '../commands/DataTransfer'
+import { IDeserializedCommand } from '../commands/CommandBase'
+import { DataTransfer, ProgressTransferResult, DataTransferState } from './dataTransfer'
+
+// TODO - this should be reimplemented on top of a generic DataTransferDownloadBuffer class
+export class DataTransferDownloadStill extends DataTransfer<Buffer> {
+	#data: Buffer[] = []
+
+	constructor(public readonly stillIndex: number) {
+		super()
+	}
+
+	public async startTransfer(transferId: number): Promise<ProgressTransferResult> {
+		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<ProgressTransferResult> {
+		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: [] }
+	}
+}

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<Buffer> {
+		const transfer = new DataTransferDownloadStill(index)
+
+		return this.#stillsLock.enqueue(transfer)
+	}
+
 	public async uploadStill(index: number, data: UploadBufferInfo, name: string, description: string): Promise<void> {
 		const transfer = new DataTransferUploadStill(index, data, name, description)
 		return this.#stillsLock.enqueue(transfer)

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)
+	})
+})

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,

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
+}