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