Skip to content

Commit cbf70e7

Browse files
AirkekAirkek
authored
Add yuv444p and nv12 support in conversion shader (#11)
* add yuv444p and nv12 support in conversion shader * fix renderer selection
1 parent aa969e1 commit cbf70e7

12 files changed

Lines changed: 386 additions & 153 deletions

‎osu-replay-viewer/Config.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ public class FFmpegOptionsObject
5252
public class OutputOptionsObject
5353
{
5454
[JsonProperty("pixel_format")] public PixelFormatMode PixelFormat = PixelFormatMode.RGB;
55+
[JsonProperty("yuv_color_space")] public ColorSpaceMode ColorSpace = ColorSpaceMode.BT709;
5556
}
5657
[JsonProperty("output_options")] public OutputOptionsObject OutputOptions = new();
5758

‎osu-replay-viewer/CustomHosts/Record/EncoderBase.cs‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,15 @@ namespace osu_replay_renderer_netcore.CustomHosts.Record;
66
public enum PixelFormatMode
77
{
88
RGB,
9-
YUV420
9+
YUV420,
10+
YUV444,
11+
NV12
12+
}
13+
14+
public enum ColorSpaceMode
15+
{
16+
BT601,
17+
BT709
1018
}
1119

1220
public struct EncoderConfig
@@ -20,6 +28,7 @@ public struct EncoderConfig
2028
public string FFmpegPath;
2129
public string FFmpegExec;
2230
public PixelFormatMode PixelFormat;
31+
public ColorSpaceMode ColorSpace;
2332
}
2433

2534
public abstract class EncoderBase

‎osu-replay-viewer/CustomHosts/Record/ExternalFFmpegEncoder.cs‎

Lines changed: 27 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
using System;
1+
using System;
22
using System.Diagnostics;
33
using System.IO;
44

@@ -12,16 +12,30 @@ private string FFmpegArguments
1212
{
1313
get
1414
{
15-
string pixFmt = "rgb24";
16-
string filters = "-vf \"vflip\"";
17-
18-
string colorFlags = "";
19-
if (Config.PixelFormat == PixelFormatMode.YUV420)
15+
string pixFmt = Config.PixelFormat switch
2016
{
21-
pixFmt = "yuv420p";
22-
filters = ""; // Shader handles flip
23-
colorFlags = "-colorspace bt709 -color_primaries bt709 -color_trc bt709 -color_range pc";
24-
}
17+
PixelFormatMode.YUV420 => "yuv420p",
18+
PixelFormatMode.YUV444 => "yuv444p",
19+
PixelFormatMode.NV12 => "nv12",
20+
_ => "rgb24"
21+
};
22+
23+
string filters = Config.PixelFormat == PixelFormatMode.RGB ? "-vf \"vflip\"" : "";
24+
25+
string colorFlags = Config.PixelFormat != PixelFormatMode.RGB ? Config.ColorSpace switch
26+
{
27+
ColorSpaceMode.BT601 => "-colorspace bt470bg -color_primaries bt470bg -color_trc gamma22 -color_range pc",
28+
ColorSpaceMode.BT709 => "-colorspace bt709 -color_primaries bt709 -color_trc bt709 -color_range pc",
29+
_ => ""
30+
} : "";
31+
32+
string outputPixFmt = Config.PixelFormat switch
33+
{
34+
PixelFormatMode.YUV420 => "yuv420p",
35+
PixelFormatMode.YUV444 => "yuv444p",
36+
PixelFormatMode.NV12 => "nv12",
37+
_ => "yuv420p" // RGB input gets converted to yuv420p by FFmpeg
38+
};
2539

2640
var inputParameters = $"-y -f rawvideo -pix_fmt {pixFmt} -s {Config.Resolution.Width}x{Config.Resolution.Height} -r {Config.FPS} -i pipe:";
2741

@@ -40,14 +54,14 @@ private string FFmpegArguments
4054
encoderSpecific = "-crf 21";
4155
break;
4256
}
43-
44-
var outputParameters = $"-c:v {Config.Encoder} {filters} {encoderSpecific} {colorFlags} -pix_fmt yuv420p -preset {Config.Preset} {Config.OutputPath}";
57+
58+
var outputParameters = $"-c:v {Config.Encoder} {filters} {encoderSpecific} {colorFlags} -pix_fmt {outputPixFmt} -preset {Config.Preset} {Config.OutputPath}";
4559
return inputParameters + (string.IsNullOrWhiteSpace(inputEffect)? (" " + inputEffect) : "") + " " + outputParameters;
4660
}
4761
}
4862

4963
public override bool CanWrite => InputStream is not null && InputStream.CanWrite;
50-
64+
5165
public ExternalFFmpegEncoder(EncoderConfig config) : base(config) { }
5266

5367
protected override void _writeFrameInternal(ReadOnlySpan<byte> frame)

‎osu-replay-viewer/CustomHosts/Record/FFmpegAutoGenEncoder.cs‎

Lines changed: 135 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ protected override void _startInternal()
3434
int bufferSize = Config.Resolution.Width * Config.Resolution.Height * 3;
3535
_pixelBuffer = ArrayPool<byte>.Shared.Rent(bufferSize);
3636
}
37-
37+
3838
ffmpeg.avformat_network_init();
3939

4040
// Allocate output format context
@@ -62,17 +62,29 @@ protected override void _startInternal()
6262
_codecContext->height = Config.Resolution.Height;
6363
_codecContext->time_base = new AVRational { num = 1, den = Config.FPS };
6464
_codecContext->framerate = new AVRational { num = Config.FPS, den = 1 };
65-
_codecContext->pix_fmt = AVPixelFormat.AV_PIX_FMT_YUV420P;
65+
_codecContext->pix_fmt = GetAVPixelFormat(Config.PixelFormat);
6666
_codecContext->codec_type = AVMediaType.AVMEDIA_TYPE_VIDEO;
6767

68-
if (Config.PixelFormat == PixelFormatMode.YUV420)
68+
// Set color metadata for YUV formats
69+
if (Config.PixelFormat != PixelFormatMode.RGB)
6970
{
70-
_codecContext->colorspace = AVColorSpace.AVCOL_SPC_BT709;
71-
_codecContext->color_primaries = AVColorPrimaries.AVCOL_PRI_BT709;
72-
_codecContext->color_trc = AVColorTransferCharacteristic.AVCOL_TRC_BT709;
71+
if (Config.ColorSpace == ColorSpaceMode.BT601)
72+
{
73+
_codecContext->colorspace = AVColorSpace.AVCOL_SPC_BT470BG;
74+
_codecContext->color_primaries = AVColorPrimaries.AVCOL_PRI_BT470BG;
75+
_codecContext->color_trc = AVColorTransferCharacteristic.AVCOL_TRC_GAMMA22;
76+
}
77+
else // BT709
78+
{
79+
_codecContext->colorspace = AVColorSpace.AVCOL_SPC_BT709;
80+
_codecContext->color_primaries = AVColorPrimaries.AVCOL_PRI_BT709;
81+
_codecContext->color_trc = AVColorTransferCharacteristic.AVCOL_TRC_BT709;
82+
}
7383
_codecContext->color_range = AVColorRange.AVCOL_RANGE_JPEG;
7484
}
7585

86+
Console.WriteLine($"[FFmpegAutoGenEncoder] Pixel format: {Config.PixelFormat} -> {_codecContext->pix_fmt}, Color space: {Config.ColorSpace}");
87+
7688
// Set encoder options
7789
var dict = new Dictionary<string, string>();
7890
dict["preset"] = Config.Preset;
@@ -128,6 +140,14 @@ protected override void _startInternal()
128140
ffmpeg.avformat_write_header(_formatContext, null);
129141
}
130142

143+
private AVPixelFormat GetAVPixelFormat(PixelFormatMode mode) => mode switch
144+
{
145+
PixelFormatMode.YUV420 => AVPixelFormat.AV_PIX_FMT_YUV420P,
146+
PixelFormatMode.YUV444 => AVPixelFormat.AV_PIX_FMT_YUV444P,
147+
PixelFormatMode.NV12 => AVPixelFormat.AV_PIX_FMT_NV12,
148+
_ => AVPixelFormat.AV_PIX_FMT_YUV420P
149+
};
150+
131151
protected override void _writeFrameInternal(ReadOnlySpan<byte> frame)
132152
{
133153
fixed (byte* framePtr = frame)
@@ -139,7 +159,7 @@ protected override void _writeFrameInternal(ReadOnlySpan<byte> frame)
139159
// For some reason sws_scale crashes with ACCESS_VIOLATION when passing mapped PBO pointer :(
140160
// TODO: find a way to avoid copying this shit
141161
Buffer.MemoryCopy(framePtr, srcPtr, _pixelBuffer.Length, frame.Length);
142-
162+
143163
int srcStride = Config.Resolution.Width * 3;
144164
byte*[] srcData = { srcPtr + (Config.Resolution.Height - 1) * srcStride, null, null, null };
145165
int[] srcStrideArray = { -srcStride, 0, 0, 0 };
@@ -149,41 +169,20 @@ protected override void _writeFrameInternal(ReadOnlySpan<byte> frame)
149169
_frame->data, _frame->linesize);
150170
}
151171
}
152-
else
172+
else if (Config.PixelFormat == PixelFormatMode.YUV420)
153173
{
154174
// YUV420P input (already flipped by shader)
155-
byte* srcPtr = framePtr;
156-
int width = Config.Resolution.Width;
157-
int height = Config.Resolution.Height;
158-
int ySize = width * height;
159-
int uvSize = width * height / 4;
160-
161-
// Y Plane
162-
byte* ySrc = srcPtr;
163-
byte* yDst = _frame->data[0];
164-
int yStride = _frame->linesize[0];
165-
for (int i = 0; i < height; i++)
166-
{
167-
Buffer.MemoryCopy(ySrc + i * width, yDst + i * yStride, yStride, width);
168-
}
169-
170-
// U Plane
171-
byte* uSrc = srcPtr + ySize;
172-
byte* uDst = _frame->data[1];
173-
int uStride = _frame->linesize[1];
174-
for (int i = 0; i < height / 2; i++)
175-
{
176-
Buffer.MemoryCopy(uSrc + i * (width / 2), uDst + i * uStride, uStride, width / 2);
177-
}
178-
179-
// V Plane
180-
byte* vSrc = srcPtr + ySize + uvSize;
181-
byte* vDst = _frame->data[2];
182-
int vStride = _frame->linesize[2];
183-
for (int i = 0; i < height / 2; i++)
184-
{
185-
Buffer.MemoryCopy(vSrc + i * (width / 2), vDst + i * vStride, vStride, width / 2);
186-
}
175+
CopyYUV420P(framePtr);
176+
}
177+
else if (Config.PixelFormat == PixelFormatMode.YUV444)
178+
{
179+
// YUV444P input
180+
CopyYUV444P(framePtr);
181+
}
182+
else if (Config.PixelFormat == PixelFormatMode.NV12)
183+
{
184+
// NV12 input
185+
CopyNV12(framePtr);
187186
}
188187

189188
_frame->pts = _pts++;
@@ -218,6 +217,101 @@ protected override void _writeFrameInternal(ReadOnlySpan<byte> frame)
218217
}
219218
}
220219

220+
private void CopyYUV420P(byte* srcPtr)
221+
{
222+
int width = Config.Resolution.Width;
223+
int height = Config.Resolution.Height;
224+
int ySize = width * height;
225+
int uvSize = width * height / 4;
226+
227+
// Y Plane
228+
byte* ySrc = srcPtr;
229+
byte* yDst = _frame->data[0];
230+
int yStride = _frame->linesize[0];
231+
for (int i = 0; i < height; i++)
232+
{
233+
Buffer.MemoryCopy(ySrc + i * width, yDst + i * yStride, yStride, width);
234+
}
235+
236+
// U Plane
237+
byte* uSrc = srcPtr + ySize;
238+
byte* uDst = _frame->data[1];
239+
int uStride = _frame->linesize[1];
240+
for (int i = 0; i < height / 2; i++)
241+
{
242+
Buffer.MemoryCopy(uSrc + i * (width / 2), uDst + i * uStride, uStride, width / 2);
243+
}
244+
245+
// V Plane
246+
byte* vSrc = srcPtr + ySize + uvSize;
247+
byte* vDst = _frame->data[2];
248+
int vStride = _frame->linesize[2];
249+
for (int i = 0; i < height / 2; i++)
250+
{
251+
Buffer.MemoryCopy(vSrc + i * (width / 2), vDst + i * vStride, vStride, width / 2);
252+
}
253+
}
254+
255+
private void CopyYUV444P(byte* srcPtr)
256+
{
257+
int width = Config.Resolution.Width;
258+
int height = Config.Resolution.Height;
259+
int planeSize = width * height;
260+
261+
// Y Plane
262+
byte* ySrc = srcPtr;
263+
byte* yDst = _frame->data[0];
264+
int yStride = _frame->linesize[0];
265+
for (int i = 0; i < height; i++)
266+
{
267+
Buffer.MemoryCopy(ySrc + i * width, yDst + i * yStride, yStride, width);
268+
}
269+
270+
// U Plane
271+
byte* uSrc = srcPtr + planeSize;
272+
byte* uDst = _frame->data[1];
273+
int uStride = _frame->linesize[1];
274+
for (int i = 0; i < height; i++)
275+
{
276+
Buffer.MemoryCopy(uSrc + i * width, uDst + i * uStride, uStride, width);
277+
}
278+
279+
// V Plane
280+
byte* vSrc = srcPtr + planeSize * 2;
281+
byte* vDst = _frame->data[2];
282+
int vStride = _frame->linesize[2];
283+
for (int i = 0; i < height; i++)
284+
{
285+
Buffer.MemoryCopy(vSrc + i * width, vDst + i * vStride, vStride, width);
286+
}
287+
}
288+
289+
private void CopyNV12(byte* srcPtr)
290+
{
291+
int width = Config.Resolution.Width;
292+
int height = Config.Resolution.Height;
293+
int ySize = width * height;
294+
int uvSize = width * height / 2; // Interleaved UV
295+
296+
// Y Plane
297+
byte* ySrc = srcPtr;
298+
byte* yDst = _frame->data[0];
299+
int yStride = _frame->linesize[0];
300+
for (int i = 0; i < height; i++)
301+
{
302+
Buffer.MemoryCopy(ySrc + i * width, yDst + i * yStride, yStride, width);
303+
}
304+
305+
// UV Plane (interleaved)
306+
byte* uvSrc = srcPtr + ySize;
307+
byte* uvDst = _frame->data[1];
308+
int uvStride = _frame->linesize[1];
309+
for (int i = 0; i < height / 2; i++)
310+
{
311+
Buffer.MemoryCopy(uvSrc + i * width, uvDst + i * uvStride, uvStride, width);
312+
}
313+
}
314+
221315
protected override void _finishInternal()
222316
{
223317
if (_pixelBuffer is not null)
@@ -290,4 +384,4 @@ private long ParseBitrate(string bitrateStr)
290384
return 10_000_000; // Fallback
291385
}
292386
}
293-
}
387+
}

‎osu-replay-viewer/CustomHosts/ReplayRecordGameHost.cs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -269,7 +269,7 @@ protected override void ChooseAndSetupRenderer()
269269

270270
if (type == GlRenderer.Auto)
271271
{
272-
if (encoder.PixelFormat == PixelFormatMode.YUV420)
272+
if (encoder.PixelFormat != PixelFormatMode.RGB)
273273
{
274274
type = GlRenderer.Legacy;
275275
}
@@ -310,7 +310,7 @@ protected override void ChooseAndSetupRenderer()
310310
}
311311

312312
SetupRendererAndWindow(rendererStr, GraphicsSurfaceType.OpenGL);
313-
wrapper = CreateWrapper(Renderer, encoder.Config.Resolution, encoder.Config.PixelFormat);
313+
wrapper = CreateWrapper(Renderer, encoder.Config.Resolution, encoder.Config.PixelFormat, encoder.Config.ColorSpace);
314314
if (wrapper is null)
315315
{
316316
Console.Error.WriteLine($"Cannot create wrapper for renderer: {Renderer.GetType()}");
@@ -320,18 +320,18 @@ protected override void ChooseAndSetupRenderer()
320320
Console.WriteLine($"Created '{type}' renderer. Type: {Renderer.GetType()}, wrapper: {wrapper.GetType()}");
321321
}
322322

323-
private static RenderWrapper CreateWrapper(IRenderer renderer, Size size, PixelFormatMode pixelFormat)
323+
private static RenderWrapper CreateWrapper(IRenderer renderer, Size size, PixelFormatMode pixelFormat, ColorSpaceMode colorSpace)
324324
{
325325
if (VeldridDeviceWrapper.IsSupported(renderer))
326326
{
327-
return new VeldridDeviceWrapper(renderer, size, pixelFormat);
327+
return new VeldridDeviceWrapper(renderer, size, pixelFormat, colorSpace);
328328
}
329329

330330
if (GLRendererWrapper.IsSupported(renderer))
331331
{
332-
return new GLRendererWrapper(renderer, size, pixelFormat);
332+
return new GLRendererWrapper(renderer, size, pixelFormat, colorSpace);
333333
}
334-
334+
335335
Console.WriteLine($"Unknown renderer: {renderer.GetType()}");
336336
throw new NotImplementedException($"Unknown renderer: {renderer.GetType()}");
337337
}

‎osu-replay-viewer/Program.cs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -256,11 +256,14 @@ static void Main(string[] args)
256256
Preset = orvConfig.FFmpegOptions.VideoEncoderPreset,
257257
Encoder = orvConfig.FFmpegOptions.VideoEncoder,
258258
Bitrate = orvConfig.FFmpegOptions.VideoEncoderBitrate,
259-
259+
260260
// External only
261261
FFmpegExec = orvConfig.FFmpegOptions.Executable,
262262
PixelFormat = orvConfig.OutputOptions.PixelFormat,
263+
ColorSpace = orvConfig.OutputOptions.ColorSpace,
263264
};
265+
266+
Console.WriteLine($"[Encoder] Pixel format: {config.PixelFormat}, Color space: {config.ColorSpace}");
264267

265268
FFmpegAudioTools.FFmpegExec = orvConfig.FFmpegOptions.Executable;
266269

0 commit comments

Comments
 (0)