forked from Qortal/Brooklyn
100 lines
2.7 KiB
C++
100 lines
2.7 KiB
C++
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;
|
|
}; |