summaryrefslogtreecommitdiff
path: root/kms++util/src/conv-yuv-planar.h
blob: b8e9247aedc38b1e654e311fa62d0132240a511a (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
#pragma once

#include <vector>

#include <kms++/framebuffer.h>
#include <kms++util/color16.h>

#include "conv-common.h"

namespace kms
{
/* YUV Planar */

template<ComponentType P1, ComponentType P2, size_t HSubUV = 1, size_t VSubUV = 1>
class YUV_Planar_Layout
	: public FormatLayout<
		PlaneLayout<uint8_t, ComponentLayout<ComponentType::Y, 8, 0>>,
		PlaneLayout<uint8_t, ComponentLayout<P1, 8, 0>>,
		PlaneLayout<uint8_t, ComponentLayout<P2, 8, 0>>
	>
{
public:
	static constexpr std::array<ComponentType, 2> uv_order = { P1, P2 };
	static constexpr size_t h_sub = HSubUV;
	static constexpr size_t v_sub = VSubUV;

	template<ComponentType P>
	static constexpr size_t find_plane()
	{
		return 1 + (std::ranges::find(uv_order, P) - uv_order.begin());
	}

	static constexpr size_t y_plane = 0; // Y is always plane 0
	static constexpr size_t cb_plane = find_plane<ComponentType::Cb>();
	static constexpr size_t cr_plane = find_plane<ComponentType::Cr>();
};

using YUV444_Layout = YUV_Planar_Layout<ComponentType::Cb, ComponentType::Cr>;
using YVU444_Layout = YUV_Planar_Layout<ComponentType::Cr, ComponentType::Cb>;
using YUV422_Layout = YUV_Planar_Layout<ComponentType::Cb, ComponentType::Cr, 2, 1>;
using YVU422_Layout = YUV_Planar_Layout<ComponentType::Cr, ComponentType::Cb, 2, 1>;
using YUV420_Layout = YUV_Planar_Layout<ComponentType::Cb, ComponentType::Cr, 2, 2>;
using YVU420_Layout = YUV_Planar_Layout<ComponentType::Cr, ComponentType::Cb, 2, 2>;

template<typename Format>
class YUVPlanarWriter
{
	using YLayout = typename Format::template plane<Format::y_plane>;
	using CbLayout = typename Format::template plane<Format::cb_plane>;
	using CrLayout = typename Format::template plane<Format::cr_plane>;

	using TY = typename YLayout::storage_type;
	using TCb = typename CbLayout::storage_type;
	using TCr = typename CrLayout::storage_type;
public:
	static void write_pattern(IFramebuffer& fb, size_t start_y, size_t end_y,
				  auto&& generate_line)
	{
		assert(start_y % Format::v_sub == 0);
		assert((end_y + 1) % Format::v_sub == 0);

		// Line buffers
		std::vector<YUV16> linebuf_storage(fb.width() * Format::v_sub);
		auto linebuf = md::mdspan(linebuf_storage.data(), Format::v_sub, fb.width());

		// Views to all planes
		auto y_buf = make_strided_fb_view<TY>(fb.map(Format::y_plane), fb.height(),
						      fb.width(), fb.stride(Format::y_plane));

		auto cb_buf = make_strided_fb_view<TCb>(fb.map(Format::cb_plane),
							fb.height() / Format::v_sub,
							fb.width() / Format::h_sub,
							fb.stride(Format::cb_plane));

		auto cr_buf = make_strided_fb_view<TCr>(fb.map(Format::cr_plane),
							fb.height() / Format::v_sub,
							fb.width() / Format::h_sub,
							fb.stride(Format::cr_plane));

		for (size_t y_src = start_y; y_src <= end_y; y_src++) {
			size_t y_offset = y_src % Format::v_sub;

			if (y_offset == 0) {
				// Fill line buffers
				for (size_t buf_y = 0; buf_y < Format::v_sub; buf_y++) {
					auto line = md::submdspan(linebuf, buf_y, md::full_extent);
					std::span<YUV16> span(line.data_handle(), line.size());
					generate_line(y_src + buf_y, span);
				}
			}

			// Write Y plane
			write_y_line(y_buf, y_src, linebuf, y_offset);

			// Write Cb/Cr planes if we're at a subsampling boundary
			if (y_offset == 0)
				write_uv_line(cb_buf, cr_buf, linebuf, y_src);
		}
	}

	static void write_lines(IFramebuffer& fb, size_t start_y, size_t end_y,
				Is2Dspan<YUV16> auto&& lines)
	{
		const size_t height = end_y - start_y + 1;

		if (lines.extent(0) < height || lines.extent(1) < fb.width())
			throw std::invalid_argument("Source line buffer too small");

		assert(start_y % Format::v_sub == 0);
		assert((end_y + 1) % Format::v_sub == 0);

		// Views to all planes
		auto y_buf = make_strided_fb_view<TY>(fb.map(Format::y_plane), fb.height(),
						      fb.width(), fb.stride(Format::y_plane));

		auto cb_buf = make_strided_fb_view<TCb>(fb.map(Format::cb_plane),
							fb.height() / Format::v_sub,
							fb.width() / Format::h_sub,
							fb.stride(Format::cb_plane));

		auto cr_buf = make_strided_fb_view<TCr>(fb.map(Format::cr_plane),
							fb.height() / Format::v_sub,
							fb.width() / Format::h_sub,
							fb.stride(Format::cr_plane));

		for (size_t y_src = start_y; y_src <= end_y; y_src++) {
			size_t y_offset = y_src % Format::v_sub;

			// Write Y plane
			write_y_line(y_buf, y_src, lines, y_offset);

			// Write Cb/Cr planes if we're at a subsampling boundary
			if (y_offset == 0)
				write_uv_line(cb_buf, cr_buf, lines, y_src);
		}
	}

private:
	template<typename YBuf>
	static void write_y_line(YBuf& y_buf, size_t y_src, auto& linebuf, size_t y_offset)
	{
		for (size_t x = 0; x < linebuf.extent(1); x++)
			y_buf(y_src, x) = YLayout::pack(linebuf(y_offset, x).y >> 8);
	}

	template<typename UVBuf>
	static void write_uv_line(UVBuf& cb_buf, UVBuf& cr_buf, auto& linebuf, size_t y_src)
	{
		const size_t uv_y = y_src / Format::v_sub;

		for (size_t x = 0; x < linebuf.extent(1); x += Format::h_sub) {
			// Average subsampled region
			uint32_t u_sum = 0, v_sum = 0;
			for (size_t y = 0; y < Format::v_sub; y++) {
				for (size_t x_off = 0; x_off < Format::h_sub; x_off++) {
					u_sum += linebuf(y, x + x_off).u;
					v_sum += linebuf(y, x + x_off).v;
				}
			}

			const size_t total_samples = Format::h_sub * Format::v_sub;
			const size_t uv_x = x / Format::h_sub;

			cb_buf(uv_y, uv_x) = CbLayout::pack(u_sum / total_samples >> 8);
			cr_buf(uv_y, uv_x) = CrLayout::pack(v_sum / total_samples >> 8);
		}
	}

public:
	static void read_lines(IFramebuffer& fb, size_t start_y, size_t end_y,
			       Is2Dspan<YUV16> auto&& dest)
	{
		const size_t height = end_y - start_y + 1;

		if (dest.extent(0) < height || dest.extent(1) < fb.width())
			throw std::invalid_argument("Destination line buffer too small");

		assert(start_y % Format::v_sub == 0);
		assert((end_y + 1) % Format::v_sub == 0);

		auto y_buf = make_strided_fb_view<const TY>(fb.map(Format::y_plane), fb.height(),
							    fb.width(), fb.stride(Format::y_plane));

		auto cb_buf = make_strided_fb_view<const TCb>(fb.map(Format::cb_plane),
							      fb.height() / Format::v_sub,
							      fb.width() / Format::h_sub,
							      fb.stride(Format::cb_plane));

		auto cr_buf = make_strided_fb_view<const TCr>(fb.map(Format::cr_plane),
							      fb.height() / Format::v_sub,
							      fb.width() / Format::h_sub,
							      fb.stride(Format::cr_plane));

		for (size_t y = start_y; y <= end_y; y++) {
			for (size_t x = 0; x < fb.width(); x++) {
				const auto y_val = YLayout::unpack(y_buf(y, x))[0];

				const size_t uv_y = y / Format::v_sub;
				const size_t uv_x = x / Format::h_sub;

				const auto u_val = CbLayout::unpack(cb_buf(uv_y, uv_x))[0];
				const auto v_val = CrLayout::unpack(cr_buf(uv_y, uv_x))[0];

				dest(y - start_y, x) = YUV16 {
					static_cast<uint16_t>(y_val << 8),
					static_cast<uint16_t>(u_val << 8),
					static_cast<uint16_t>(v_val << 8)
				};
			}
		}
	}
};

} // namespace kms