mirror of https://github.com/Qortal/Brooklyn
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
100 lines
2.7 KiB
100 lines
2.7 KiB
using namespace boost::asio; |
|
using ip::tcp; |
|
|
|
typedef boost::shared_ptr<tcp::socket> socket_ptr; |
|
|
|
#include "NvEncoder/NvEncoder.h" |
|
#include "color_conversion.h" |
|
|
|
class NV_encoding { |
|
public: |
|
void load(int width, int height, socket_ptr sock, UINT monitorID) { |
|
NVENCSTATUS nvStatus = NV_ENC_SUCCESS; |
|
|
|
this->sock = sock; |
|
this->width = width; |
|
this->height = height; |
|
cNvEncoder = new CNvEncoder(); |
|
cNvEncoder->InitCuda(); |
|
nvStatus = cNvEncoder->Initialize(NV_ENC_DEVICE_TYPE_CUDA); |
|
nvStatus = cNvEncoder->CreateEncoder(width, height); |
|
nvStatus = cNvEncoder->AllocateIOBuffers(width, height, false); |
|
|
|
dataPacket = new DataPacket(); |
|
dataPacket->data = new uint8_t[width*height]; |
|
|
|
yuv[0] = new uint8_t[width*height]; |
|
yuv[1] = new uint8_t[width*height / 4]; |
|
yuv[2] = new uint8_t[width*height / 4]; |
|
|
|
// Init avi file |
|
//char buffer[255]; |
|
//sprintf(buffer, "C:\\Monitor%d.avi", monitorID); |
|
//ofs.open(buffer, std::ofstream::out | std::ofstream::binary); |
|
|
|
} |
|
void write(int width, int height, RGBQUAD *pPixels) { |
|
|
|
bool rc = RGB_to_YV12(width, height, pPixels, yuv[0], yuv[1], yuv[2]); |
|
|
|
if (!rc){ |
|
// The Cuda function RGB_to_YV12 failed, do CPU conversion |
|
for (int y = 0; y < height; y++) { |
|
for (int x = 0; x < width; x++) { |
|
|
|
RGBQUAD px = pPixels[y*width+x]; |
|
int Y = ( ( 66 * px.rgbRed + 129 * px.rgbGreen + 25 * px.rgbBlue + 128) >> 8) + 16; |
|
int U = ( ( -38 * px.rgbRed - 74 * px.rgbGreen + 112 * px.rgbBlue + 128) >> 8) + 128; |
|
int V = ( ( 112 * px.rgbRed - 94 * px.rgbGreen - 18 * px.rgbBlue + 128) >> 8) + 128; |
|
|
|
yuv[0][y * width + x] = Y; |
|
yuv[1][(y >> 1) * (width >> 1) + (x >> 1)] = U; |
|
yuv[2][(y >> 1) * (width >> 1) + (x >> 1)] = V; |
|
} |
|
} |
|
} |
|
|
|
EncodeFrameConfig stEncodeFrame; |
|
memset(&stEncodeFrame, 0, sizeof(stEncodeFrame)); |
|
|
|
stEncodeFrame.yuv[0] = yuv[0]; |
|
stEncodeFrame.yuv[1] = yuv[1]; |
|
stEncodeFrame.yuv[2] = yuv[2]; |
|
|
|
stEncodeFrame.stride[0] = width; |
|
stEncodeFrame.stride[1] = width/2; |
|
stEncodeFrame.stride[2] = width/2; |
|
stEncodeFrame.width = width; |
|
stEncodeFrame.height = height; |
|
|
|
cNvEncoder->EncodeFrame(&stEncodeFrame, dataPacket, false, width, height); |
|
if (dataPacket->size > 0) { |
|
printf("Write frame (size=%5d)\n", dataPacket->size); |
|
|
|
//ofs.write((char*)dataPacket->data, dataPacket->size); |
|
|
|
boost::asio::write(*sock, buffer((char*)dataPacket->data, dataPacket->size)); |
|
} |
|
} |
|
void close () { |
|
delete cNvEncoder; |
|
delete dataPacket->data; |
|
delete dataPacket; |
|
for (int i = 0; i < 3; i ++) |
|
{ |
|
if (yuv[i]) |
|
{ |
|
delete [] yuv[i]; |
|
} |
|
} |
|
} |
|
private: |
|
int width; |
|
int height; |
|
socket_ptr sock; |
|
uint8_t *yuv[3]; |
|
CNvEncoder* cNvEncoder; |
|
DataPacket* dataPacket; |
|
|
|
//std::ofstream ofs; |
|
}; |