1 #pragma once
2 
3 #include <obs-module.h>
4 
5 #define WIN32_MEAN_AND_LEAN
6 #include <Windows.h>
7 #undef WIN32_MEAN_AND_LEAN
8 
9 #include <mfapi.h>
10 #include <mfidl.h>
11 
12 #include <wmcodecdsp.h>
13 
14 #include <vector>
15 #include <queue>
16 #include <memory>
17 #include <atomic>
18 
19 #include <util/windows/ComPtr.hpp>
20 
21 #include "mf-encoder-descriptor.hpp"
22 #include "mf-common.hpp"
23 
24 namespace MF {
25 	enum H264Profile {
26 		H264ProfileBaseline,
27 		H264ProfileMain,
28 		H264ProfileHigh
29 	};
30 
31 	enum H264RateControl {
32 		H264RateControlCBR,
33 		H264RateControlConstrainedVBR,
34 		H264RateControlVBR,
35 		H264RateControlCQP
36 	};
37 
38 	struct H264QP {
39 		UINT16 defaultQp;
40 		UINT16 i;
41 		UINT16 p;
42 		UINT16 b;
43 
PackMF::H264QP44 		UINT64 Pack(bool packDefault) {
45 			int shift = packDefault ? 0 : 16;
46 			UINT64 packedQp;
47 			if (packDefault)
48 				packedQp = defaultQp;
49 
50 			packedQp |= i << shift;
51 			shift += 16;
52 			packedQp |= p << shift;
53 			shift += 16;
54 			packedQp |= b << shift;
55 
56 			return packedQp;
57 		}
58 	};
59 
60 	enum H264EntropyEncoding {
61 		H264EntropyEncodingCABLC,
62 		H264EntropyEncodingCABAC
63 	};
64 
65 	struct H264Frame {
66 	public:
H264FrameMF::H264Frame67 		H264Frame(bool keyframe, UINT64 pts, UINT64 dts,
68 				std::unique_ptr<std::vector<uint8_t>> data)
69 			: keyframe(keyframe), pts(pts), dts(dts),
70 			  data(std::move(data))
71 		{}
KeyframeMF::H264Frame72 		bool Keyframe() { return keyframe; }
DataMF::H264Frame73 		BYTE *Data() { return data.get()->data(); }
DataLengthMF::H264Frame74 		DWORD DataLength() { return (DWORD)data.get()->size(); }
PtsMF::H264Frame75 		INT64 Pts() { return pts; }
DtsMF::H264Frame76 		INT64 Dts() { return dts; }
77 
78 	private:
79 		H264Frame(H264Frame const&) = delete;
80 		H264Frame& operator=(H264Frame const&) = delete;
81 	private:
82 		bool keyframe;
83 		INT64 pts;
84 		INT64 dts;
85 		std::unique_ptr<std::vector<uint8_t>> data;
86 	};
87 
88 	class H264Encoder {
89 	public:
90 		H264Encoder(const obs_encoder_t *encoder,
91 			std::shared_ptr<EncoderDescriptor> descriptor,
92 			UINT32 width,
93 			UINT32 height,
94 			UINT32 framerateNum,
95 			UINT32 framerateDen,
96 			H264Profile profile,
97 			UINT32 bitrate);
98 
99 		~H264Encoder();
100 
101 		bool Initialize(std::function<bool(void)> func);
102 		bool ProcessInput(UINT8 **data, UINT32 *linesize, UINT64 pts,
103 			Status *status);
104 		bool ProcessOutput(UINT8 **data, UINT32 *dataLength,
105 			UINT64 *pts, UINT64 *dts, bool *keyframe,
106 			Status *status);
107 		bool ExtraData(UINT8 **data, UINT32 *dataLength);
108 
ObsEncoder()109 		const obs_encoder_t *ObsEncoder() { return encoder; }
110 
111 	public:
112 		bool SetBitrate(UINT32 bitrate);
113 		bool SetQP(H264QP &qp);
114 		bool SetMaxBitrate(UINT32 maxBitrate);
115 		bool SetRateControl(H264RateControl rateControl);
116 		bool SetKeyframeInterval(UINT32 seconds);
117 		bool SetLowLatency(bool lowLatency);
118 		bool SetBufferSize(UINT32 bufferSize);
119 		bool SetBFrameCount(UINT32 bFrames);
120 		bool SetEntropyEncoding(H264EntropyEncoding entropyEncoding);
121 		bool SetMinQP(UINT32 minQp);
122 		bool SetMaxQP(UINT32 maxQp);
123 
124 	private:
125 		H264Encoder(H264Encoder const&) = delete;
126 		H264Encoder& operator=(H264Encoder const&) = delete;
127 
128 	private:
129 		HRESULT InitializeEventGenerator();
130 		HRESULT InitializeExtraData();
131 		HRESULT CreateMediaTypes(ComPtr<IMFMediaType> &inputType,
132 			ComPtr<IMFMediaType> &outputType);
133 		HRESULT EnsureCapacity(ComPtr<IMFSample> &sample, DWORD length);
134 		HRESULT CreateEmptySample(ComPtr<IMFSample> &sample,
135 			ComPtr<IMFMediaBuffer> &buffer, DWORD length);
136 
137 		HRESULT ProcessInput(ComPtr<IMFSample> &sample);
138 		HRESULT ProcessOutput();
139 
140 		HRESULT DrainEvent(bool block);
141 		HRESULT DrainEvents();
142 	private:
143 		const obs_encoder_t *encoder;
144 		std::shared_ptr<EncoderDescriptor> descriptor;
145 		const UINT32 width;
146 		const UINT32 height;
147 		const UINT32 framerateNum;
148 		const UINT32 framerateDen;
149 		const UINT32 initialBitrate;
150 		const H264Profile profile;
151 
152 		bool createOutputSample;
153 		ComPtr<IMFTransform> transform;
154 		ComPtr<ICodecAPI> codecApi;
155 
156 		std::vector<BYTE> extraData;
157 
158 		// The frame returned by ProcessOutput
159 		// Valid until the next call to ProcessOutput
160 		std::unique_ptr<H264Frame> activeFrame;
161 
162 		// Queued input samples that the encoder was not ready
163 		// to process
164 		std::queue<ComPtr<IMFSample>> inputSamples;
165 
166 		// Queued output samples that have not been returned from
167 		// ProcessOutput yet
168 		std::queue<std::unique_ptr<H264Frame>> encodedFrames;
169 
170 		ComPtr<IMFMediaEventGenerator> eventGenerator;
171 		std::atomic<UINT32> inputRequests = 0;
172 		std::atomic<UINT32> outputRequests = 0;
173 		std::atomic<UINT32> pendingRequests = 0;
174 	};
175 }
176