Skip to content

Commit

Permalink
Decompress JPEG lossy to YUV instead of BGRA
Browse files Browse the repository at this point in the history 
This commit fixes several situations where the receiver would treat color ranges wrong when JPEG lossy was used, colors were displayed wrong e.g. for BGRA full and partial and YUV (NV12, I420, I444) partial sender color settings.
  • Loading branch information
@YorVeX
YorVeX committed May 29, 2023
1 parent 7fc6035 commit 03cf2ae
Show file tree
Hide file tree
Showing 5 changed files with 80 additions and 32 deletions.
5 changes: 3 additions & 2 deletions Beam.cs
View file
Expand Up @@ -15,7 +15,8 @@ public enum CompressionTypes : int
Qoi = 1,
Lz4 = 2,
QoiLz4 = 3,
Jpeg = 4,
JpegLossy = 4,
JpegLossless = 5,
}

#region helper methods
Expand Down Expand Up @@ -43,7 +44,7 @@ public static SequencePosition GetTimestamp(ReadOnlySequence<byte> sequence, out
return reader.Position;
}

public static unsafe uint[] GetYuvPlaneSizes(video_format format, uint width, uint height)
public static uint[] GetYuvPlaneSizes(video_format format, uint width, uint height)
{
uint halfHeight;
uint halfwidth;
Expand Down
47 changes: 27 additions & 20 deletions BeamReceiver.cs
View file
Expand Up @@ -10,7 +10,6 @@
using System.Runtime.InteropServices;
using LibJpegTurbo;
using K4os.Compression.LZ4;
using ObsInterop;

namespace xObsBeam;

Expand Down Expand Up @@ -276,27 +275,19 @@ private unsafe void TurboJpegDecompressDestroy()
}
}

private static unsafe uint GetRawVideoDataSize(Beam.VideoHeader videoHeader, out uint[] videoPlaneSizes)
private static unsafe uint GetRawVideoDataSize(Beam.VideoHeader videoHeader)
{
uint rawVideoDataSize = 0;
// get the plane sizes for the current frame format and size
videoPlaneSizes = Beam.GetPlaneSizes(videoHeader.Format, videoHeader.Height, videoHeader.Linesize);
var videoPlaneSizes = Beam.GetPlaneSizes(videoHeader.Format, videoHeader.Height, videoHeader.Linesize);
if (videoPlaneSizes.Length == 0) // unsupported format
return rawVideoDataSize;

for (int planeIndex = 0; planeIndex < videoPlaneSizes.Length; planeIndex++)
rawVideoDataSize += videoPlaneSizes[planeIndex];

return rawVideoDataSize;
}
private unsafe void ChangeHeaderToBgra(ref Beam.VideoHeader videoHeader)
{
videoHeader.Format = video_format.VIDEO_FORMAT_BGRA;
fixed (uint* linesize = videoHeader.Linesize)
{
new Span<uint>(linesize, Beam.VideoHeader.MAX_AV_PLANES).Clear(); // BGRA means only one plane, don't use the original YUV linesizes
linesize[0] = videoHeader.Width * 4; // BGRA has 4 bytes per pixel
}
}

#endregion unsafe helper functions

Expand Down Expand Up @@ -331,6 +322,7 @@ private async Task ProcessDataLoopAsync(Socket? socket, NamedPipeClientStream? p

int videoHeaderSize = Beam.VideoHeader.VideoHeaderDataSize;
uint rawVideoDataSize = 0;
uint[] videoPlaneSizes = new uint[Beam.VideoHeader.MAX_AV_PLANES];

uint fps = 30;
uint logCycle = 0;
Expand Down Expand Up @@ -425,7 +417,24 @@ private async Task ProcessDataLoopAsync(Socket? socket, NamedPipeClientStream? p
if (sizeChanged || firstFrame) // re-allocate the arrays matching the new necessary size
{
firstFrame = false;
rawVideoDataSize = GetRawVideoDataSize(videoHeader, out _);

if (videoHeader.Compression == Beam.CompressionTypes.JpegLossy)
{
if (EncoderSupport.LibJpegTurboV3)
rawVideoDataSize = (uint)TurboJpeg.tj3YUVBufSize((int)videoHeader.Width, 1, (int)videoHeader.Height, (int)TJSAMP.TJSAMP_420);
else if (EncoderSupport.LibJpegTurbo)
rawVideoDataSize = (uint)TurboJpeg.tjBufSizeYUV2((int)videoHeader.Width, 1, (int)videoHeader.Height, (int)TJSAMP.TJSAMP_420);
else
{
rawVideoDataSize = 0;
Module.Log($"Error: JPEG library is not available, cannot decompress received video data!", ObsLogLevel.Error);
break;
}
EncoderSupport.GetJpegPlaneSizes((int)videoHeader.Width, (int)videoHeader.Height, out videoPlaneSizes, out _);
}
else
rawVideoDataSize = GetRawVideoDataSize(videoHeader);

maxVideoDataSize = (int)(videoHeader.Width * videoHeader.Height * 4);
receivedFrameData = new byte[rawVideoDataSize];
lz4DecompressBuffer = new byte[maxVideoDataSize];
Expand Down Expand Up @@ -460,14 +469,12 @@ private async Task ProcessDataLoopAsync(Socket? socket, NamedPipeClientStream? p
// need to decompress QOI only
else if (videoHeader.Compression == Beam.CompressionTypes.Qoi)
Qoi.Decode(receivedFrameData, videoHeader.DataSize, rawDataBuffer, maxVideoDataSize);
// need to decompress JPEG only
else if (videoHeader.Compression == Beam.CompressionTypes.Jpeg)
{
//TODO: output to YUV instead of BGRA by default, the necessary turboJpegDecompressToYuv() function is already there but hasn't been tested yet
// this would save libjpeg-turbo from having to do a conversion from native JPEG YUV to BGRA, but also add significant code complexity since the plane sizes need to be determined, the question is whether that's actually worth it
// need to decompress JPEG lossless only
else if (videoHeader.Compression == Beam.CompressionTypes.JpegLossless)
TurboJpegDecompressToBgra(receivedFrameData, maxVideoDataSize, rawDataBuffer, (int)videoHeader.Width, (int)videoHeader.Height);
ChangeHeaderToBgra(ref videoHeader); // we decompressed to BGRA
}
// need to decompress JPEG lossy only
else if (videoHeader.Compression == Beam.CompressionTypes.JpegLossy)
TurboJpegDecompressToYuv(receivedFrameData, maxVideoDataSize, rawDataBuffer, videoPlaneSizes, (int)videoHeader.Width, (int)videoHeader.Height);
}

// process the frame
Expand Down
14 changes: 9 additions & 5 deletions BeamSender.cs
View file
Expand Up @@ -91,6 +91,9 @@ public unsafe bool SetVideoParameters(video_output_info* info, video_format conv
return false;
}

for (int i = 0; i < Beam.VideoHeader.MAX_AV_PLANES; i++)
Module.Log("SetVideoParameters(): linesize[" + i + "] = " + linesize[i], ObsLogLevel.Debug);

var pointerOffset = (IntPtr)data.e0;
for (int planeIndex = 0; planeIndex < _videoPlaneSizes.Length; planeIndex++)
{
Expand All @@ -101,7 +104,7 @@ public unsafe bool SetVideoParameters(video_output_info* info, video_format conv
if (pointerOffset != (IntPtr)data[planeIndex])
{
// either the GetPlaneSizes() returned wrong information or the video data plane pointers are not contiguous in memory (which we currently rely on)
Module.Log($"Video data plane pointer for plane {planeIndex} of format {info->format} has a difference of {pointerOffset - (IntPtr)data[planeIndex]}.", ObsLogLevel.Error);
Module.Log($"Video data plane pointer for plane {planeIndex} of format {info->format} has a difference of {pointerOffset - (IntPtr)data[planeIndex]}.", ObsLogLevel.Warning);
//BUG: this is currently happening for odd resolutions like 1279x719 on YUV formats, because padding is not properly handled by GetPlaneSizes(), leading to image distortion
}
pointerOffset += (int)_videoPlaneSizes[planeIndex];
Expand Down Expand Up @@ -163,12 +166,12 @@ public unsafe bool SetVideoParameters(video_output_info* info, video_format conv
}

var videoBandwidthMbps = (((Beam.VideoHeader.VideoHeaderDataSize + _videoDataSize) * (info->fps_num / info->fps_den)) / 1024 / 1024) * 8;
if (!_qoiCompression && !_lz4Compression)
if (!_qoiCompression && !_lz4Compression && !_jpegCompression)
Module.Log($"Video output feed initialized, theoretical uncompressed net bandwidth demand is {videoBandwidthMbps} Mpbs.", ObsLogLevel.Info);
else
{
string lz4WithLevelString = $"LZ4 ({SettingsDialog.Lz4CompressionLevel})";
Module.Log($"Video output feed initialized with {(_qoiCompression ? (_lz4Compression ? "QOI + " + lz4WithLevelString : "QOI") : lz4WithLevelString)} compression. Sync to render thread: {_compressionThreadingSync}. Theoretical uncompressed net bandwidth demand would be {videoBandwidthMbps} Mpbs.", ObsLogLevel.Info);
Module.Log($"Video output feed initialized with {(_jpegCompression ? "JPEG" : (_qoiCompression ? (_lz4Compression ? "QOI + " + lz4WithLevelString : "QOI") : lz4WithLevelString))} compression. Sync to render thread: {_compressionThreadingSync}. Theoretical uncompressed net bandwidth demand would be {videoBandwidthMbps} Mpbs.", ObsLogLevel.Info);
}

return true;
Expand Down Expand Up @@ -411,7 +414,7 @@ private unsafe void SendCompressed(ulong timestamp, Beam.VideoHeader videoHeader

if (encodedDataLength < videoHeader.DataSize) // did compression decrease the size of the data?
{
videoHeader.Compression = Beam.CompressionTypes.Jpeg;
videoHeader.Compression = (_jpegCompressionLossless ? Beam.CompressionTypes.JpegLossless : Beam.CompressionTypes.JpegLossy);
videoHeader.DataSize = encodedDataLength;
}
}
Expand Down Expand Up @@ -454,7 +457,8 @@ private unsafe void SendCompressed(ulong timestamp, Beam.VideoHeader videoHeader

switch (videoHeader.Compression)
{
case Beam.CompressionTypes.Jpeg:
case Beam.CompressionTypes.JpegLossy:
case Beam.CompressionTypes.JpegLossless:
foreach (var client in _clients.Values)
client.EnqueueVideoFrame(timestamp, videoHeader, encodedDataJpeg!);
break;
Expand Down
24 changes: 24 additions & 0 deletions EncoderSupport.cs
View file
Expand Up @@ -185,6 +185,30 @@ public static TJCS ObsToJpegColorSpace(video_format obsVideoFormat)
return (FormatIsYuv(obsVideoFormat)) ? TJCS.TJCS_YCbCr : TJCS.TJCS_RGB;
}

public static void GetJpegPlaneSizes(int width, int height, out uint[] videoPlaneSizes, out uint[] linesize)
{
videoPlaneSizes = new uint[Beam.VideoHeader.MAX_AV_PLANES];
linesize = new uint[Beam.VideoHeader.MAX_AV_PLANES];
if (LibJpegTurboV3)
{
for (int i = 0; i < videoPlaneSizes.Length; i++)
{
videoPlaneSizes[i] = (uint)TurboJpeg.tj3YUVPlaneSize(i, width, 0, height, (int)TJSAMP.TJSAMP_420);
linesize[i] = (uint)TurboJpeg.tj3YUVPlaneWidth(i, width, (int)TJSAMP.TJSAMP_420);
}
}
else if (LibJpegTurbo)
{
for (int i = 0; i < videoPlaneSizes.Length; i++)
{
videoPlaneSizes[i] = (uint)TurboJpeg.tjPlaneSizeYUV(i, width, 0, height, (int)TJSAMP.TJSAMP_420);
linesize[i] = (uint)TurboJpeg.tjPlaneWidth(i, width, (int)TJSAMP.TJSAMP_420);
}
}
else
Module.Log($"Error: JPEG library is not available, cannot get JPEG plane sizes!", ObsLogLevel.Error);
}

[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static unsafe void Nv12ToI420(byte* sourceBuffer, Span<byte> destinationBuffer, uint[] planeSizes)
{
Expand Down
22 changes: 17 additions & 5 deletions Source.cs
View file
Expand Up @@ -374,13 +374,25 @@ private unsafe void VideoFrameReceivedEventHandler(object? sender, Beam.BeamVide
context->Video->width = videoFrame.Header.Width;
context->Video->height = videoFrame.Header.Height;
context->Video->full_range = Convert.ToByte(videoFrame.Header.Range == video_range_type.VIDEO_RANGE_FULL);
for (int i = 0; i < Beam.VideoHeader.MAX_AV_PLANES; i++)
// get the plane sizes for the current frame format and size
if (videoFrame.Header.Compression == Beam.CompressionTypes.JpegLossy)
{
context->Video->linesize[i] = videoFrame.Header.Linesize[i];
Module.Log("VideoFrameReceivedEventHandler(): linesize[" + i + "] = " + context->Video->linesize[i], ObsLogLevel.Debug);
EncoderSupport.GetJpegPlaneSizes((int)context->Video->width, (int)context->Video->height, out _videoPlaneSizes, out var jpeglineSize);
for (int i = 0; i < Beam.VideoHeader.MAX_AV_PLANES; i++)
{
context->Video->linesize[i] = jpeglineSize[i];
Module.Log("VideoFrameReceivedEventHandler(): linesize[" + i + "] = " + context->Video->linesize[i], ObsLogLevel.Debug);
}
}
else
{
_videoPlaneSizes = Beam.GetPlaneSizes(context->Video->format, context->Video->height, context->Video->linesize);
for (int i = 0; i < Beam.VideoHeader.MAX_AV_PLANES; i++)
{
context->Video->linesize[i] = videoFrame.Header.Linesize[i];
Module.Log("VideoFrameReceivedEventHandler(): linesize[" + i + "] = " + context->Video->linesize[i], ObsLogLevel.Debug);
}
}
// get the plane sizes for the current frame format and size
_videoPlaneSizes = Beam.GetPlaneSizes(context->Video->format, context->Video->height, context->Video->linesize);
ObsVideo.video_format_get_parameters_for_format(videoFrame.Header.Colorspace, videoFrame.Header.Range, videoFrame.Header.Format, context->Video->color_matrix, context->Video->color_range_min, context->Video->color_range_max);
Module.Log("VideoFrameReceivedEventHandler(): reinitialized", ObsLogLevel.Debug);
}
Expand Down

0 comments on commit 03cf2ae

Please sign in to comment.