diff --git a/Dockerfile b/Dockerfile index 606c515..7650e39 100644 --- a/Dockerfile +++ b/Dockerfile @@ -5,6 +5,7 @@ ARG BUILD_THREAD=5 RUN apt update RUN apt install -y build-essential libboost-all-dev libsqlite3-dev libasio-dev nasm +RUN apt install -y ffmpeg libswscale-dev libavutil-dev libavcodec-dev libavdevice-dev libavfilter-dev libavformat-dev libavutil-dev libpostproc-dev libswresample-dev WORKDIR /app @@ -13,7 +14,7 @@ COPY . . RUN cd x264; make -j $BUILD_THREAD; cd .. RUN cd zstd; make -j $BUILD_THREAD; cd .. -RUN make +RUN make all EXPOSE ${CDN_PORT} diff --git a/Makefile b/Makefile index 24288b6..5e1fbe9 100644 --- a/Makefile +++ b/Makefile @@ -1,2 +1,7 @@ -app: - g++ -std=c++17 -O3 -march=native src/app.cpp ./zstd/lib/libzstd.so ./x264/libx264.a -lpthread -lsqlite3 -o nTracer_cdn -Wl,--no-as-needed -ldl \ No newline at end of file +CXXFILES = src/app.cpp +INCLIBS = ./zstd/lib/libzstd.so ./x264/libx264.a +LIBS = -lpthread -lsqlite3 -lavcodec -lavformat -lavutil -lswscale +CXXFLAGS = -O3 -std=c++17 -march=native -o nTracer_cdn -Wl,--no-as-needed -ldl + +all: + $(CXX) $(CXXFILES) $(INCLIBS) $(LIBS) $(CXXFLAGS) \ No newline at end of file diff --git a/src/reader.hpp b/src/reader.hpp index b321c7c..85df0e6 100644 --- a/src/reader.hpp +++ b/src/reader.hpp @@ -266,7 +266,7 @@ class packed_reader case 2: // Decompress with vidlib - read_decomp_buffer_pt = decode_stack(chunkx, chunky, chunkz, read_buffer, sel->size); + read_decomp_buffer_pt = decode_stack_2(chunkx, chunky, chunkz, read_buffer, sel->size); read_decomp_buffer = (char *)pixtype_to_uint16(read_decomp_buffer_pt, chunkx * chunky * chunkz); free(read_decomp_buffer_pt); break; diff --git a/src/vidlib.hpp b/src/vidlib.hpp index 20b95b6..8b3fd05 100644 --- a/src/vidlib.hpp +++ b/src/vidlib.hpp @@ -4,17 +4,30 @@ #include #include +#include #include #include +#include +#include #include "x264.h" +extern "C" +{ +#include +#include +#include +#include +#include +#include +#include +} + #define IMAGE_GAIN 10 std::mutex pthread_mutex; -std::string ffmpeg_location = "/home/loganaw/test/nTracer2_cdn3/ffmpeg-6.0-amd64-static/ffmpeg"; -// std::string ffmpeg_location = "/Users/loganaw/Downloads/ffmpeg"; +std::string ffmpeg_location = "./ffmpeg"; std::string null_redirect = ""; // "2>/dev/null"; std::string encoder_name = "libx264"; std::string other_encode_settings = ""; @@ -55,10 +68,10 @@ void read_threaded(void **buffer_ret, size_t *buffer_size, FILE *file, size_t si buffer_ret[0] = buffer; } -pixtype *decode_stack(size_t w, size_t h, size_t t, void *buffer, size_t buffer_size) +pixtype *decode_stack_subprocess(size_t sizex, size_t sizey, size_t sizez, void *buffer, size_t buffer_size) { - const size_t frame_size = w * h * sizeof(pixtype); - const size_t stack_size = frame_size * t; + const size_t frame_size = sizex * sizey * sizeof(pixtype); + const size_t stack_size = frame_size * sizez; std::string cmd = ffmpeg_location + " " + decode_params; @@ -68,46 +81,39 @@ pixtype *decode_stack(size_t w, size_t h, size_t t, void *buffer, size_t buffer_ subprocess::output{subprocess::PIPE}, subprocess::input{subprocess::PIPE}, subprocess::error{subprocess::PIPE}, - subprocess::bufsize{(int) 256 * 256 * 256 * 2}); - - //std::cerr << "GOT " << buffer_size << std::endl; - - //{ - // std::ofstream file("test.h256", std::ios::binary); - // file.write((char *)buffer, buffer_size); - // file.close(); - //} + subprocess::bufsize{(int)256 * 256 * 256 * 2}); -/* - const size_t send_chunk_size = 1000; - for(size_t i = 0; i < buffer_size; i+=send_chunk_size) { - char * buffer_offset = (char *) buffer; - buffer_offset += i; + // size_t s = p->send((char *)buffer, buffer_size); + // std::cerr << "SENT " << s << std::endl; - std::cerr << "PREPPING " << i << std::endl; + // std::pair res = p->communicate((char *) buffer, buffer_size); + // std::vector buf = res.first.buf; - size_t tosend = std::min(send_chunk_size, buffer_size - i); - - size_t s = p->send(buffer_offset, tosend); - - std::cerr << "SENTT " << s << std::endl; - } - */ + auto buf = p->communicate((char *)buffer, buffer_size).first.buf; + size_t buf_size = buf.size(); - //size_t s = p->send((char *)buffer, buffer_size); - //std::cerr << "SENT " << s << std::endl; + // std::cerr << "Got " << buf_size << std::endl; - auto buf = p->communicate((char *) buffer, buffer_size).first.buf; + char *out = (char *)calloc(stack_size, sizeof(char)); - //auto buf = p->communicate().first.buf; + // Data stored in ZXY, remap to XYZ + for (size_t x = 0; x < sizex; x++) + { + for (size_t y = 0; y < sizey; y++) + { + for (size_t z = 0; z < sizez; z++) + { + const size_t in_offset = (z * (sizex * sizey)) + (x * sizey) + y; - size_t buf_size = buf.size(); - //std::cerr << "Got " << buf_size << std::endl; - void *out = malloc(buf_size); - memcpy(out, buf.data(), buf_size); + const size_t out_offset = (x * sizey * sizez) + (y * sizez) + z; - if (buf_size != stack_size) - ; + if (in_offset < buf_size) + { + out[out_offset] = buf[in_offset]; + } + } + } + } delete p; @@ -164,7 +170,7 @@ std::pair encode_stack(size_t w, size_t h, size_t t, pixtype *st subprocess::input{subprocess::PIPE}, subprocess::error{subprocess::PIPE}, subprocess::close_fds{false}, - subprocess::bufsize{(int) 256 * 256 * 256 * 2}, + subprocess::bufsize{(int)256 * 256 * 256 * 2}, subprocess::shell{false}); auto [outb, errb] = p->communicate((char *)stack, stack_size); @@ -186,13 +192,13 @@ std::pair encode_stack_AV1_encapp(size_t w, size_t h, size_t t, { std::stringstream cmd_build; cmd_build << SVT_AV1_location << " " - << "--crf 20 " - << "--lp 4 " - << "-i - " - << "-w " << w << " " - << "-h " << h << " " - << "--fps " << target_framerate << " " - << "-b -"; + << "--crf 20 " + << "--lp 4 " + << "-i - " + << "-w " << w << " " + << "-h " << h << " " + << "--fps " << target_framerate << " " + << "-b -"; // Use subprocess subprocess::Popen *p = new subprocess::Popen( @@ -201,7 +207,7 @@ std::pair encode_stack_AV1_encapp(size_t w, size_t h, size_t t, subprocess::input{subprocess::PIPE}, subprocess::error{subprocess::PIPE}, subprocess::close_fds{false}, - subprocess::bufsize{(int) (w * h * t * 2)}, + subprocess::bufsize{(int)(w * h * t * 2)}, subprocess::shell{false}); auto [outb, errb] = p->communicate((char *)stack, stack_size_in); @@ -340,7 +346,6 @@ pixtype *uint16_to_pixtype(uint16_t *buffer, size_t len) v *= IMAGE_GAIN; // Apply input gain v = sqrt(v); // sqrt image to do 16->8bit conversion - if (v <= 0) v = 0; // Clip lower if (v >= 255) @@ -442,3 +447,221 @@ uint16_t *pixtype_to_uint16_YUV420(pixtype *buffer, size_t w, size_t h, size_t t return out_buffer; } + +// Custom I/O context for reading from memory +class FFmpegMemoryBuffer +{ +public: + const uint8_t *data; + size_t size; + size_t pos; + + FFmpegMemoryBuffer(const uint8_t *buffer, size_t bufferSize) + : data(buffer), size(bufferSize), pos(0) {} +}; + +// Custom read function for memory buffer +static int memorybuffer_read_packet(void *opaque, uint8_t *buf, int buf_size) +{ + FFmpegMemoryBuffer *memBuffer = static_cast(opaque); + + // Calculate how much we can read + int bytesToRead = std::min(buf_size, static_cast(memBuffer->size - memBuffer->pos)); + + if (bytesToRead <= 0) + { + return AVERROR_EOF; + } + + // Copy data from memory buffer + memcpy(buf, memBuffer->data + memBuffer->pos, bytesToRead); + memBuffer->pos += bytesToRead; + + return bytesToRead; +} + +pixtype *decode_stack_2(size_t sizex, size_t sizey, size_t sizez, void *buffer, size_t buffer_size) +{ + // Allocate output buffer + uint8_t *out = (uint8_t *)calloc(sizex * sizey * sizez, sizeof(uint8_t)); + + if (!buffer || buffer_size == 0) + { + std::cerr << "[H264Decode] Failed to load buffer" << std::endl; + return (pixtype *)out; + } + + // Create memory buffer context + FFmpegMemoryBuffer memBuffer((const uint8_t *)buffer, buffer_size); + + // Allocate AVFormatContext + AVFormatContext *formatContext = avformat_alloc_context(); + if (!formatContext) + { + std::cerr << "[H264Decode] Could not allocate format context" << std::endl; + return (pixtype *)out; + } + + // Create custom I/O context + AVIOContext *ioContext = avio_alloc_context( + static_cast(av_malloc(4096)), // Internal buffer + 4096, // Buffer size + 0, // Write flag (0 for read-only) + &memBuffer, // Opaque pointer + memorybuffer_read_packet, // Read callback + nullptr, // Write callback (not needed) + nullptr // Seek callback (optional) + ); + + if (!ioContext) + { + std::cerr << "[H264Decode] Could not create I/O context" << std::endl; + avformat_free_context(formatContext); + return (pixtype *)out; + } + + // Assign custom I/O context to format context + formatContext->pb = ioContext; + + // Open input from memory buffer + int ret = avformat_open_input(&formatContext, nullptr, nullptr, nullptr); + if (ret < 0) + { + char errbuf[AV_ERROR_MAX_STRING_SIZE]; + av_strerror(ret, errbuf, AV_ERROR_MAX_STRING_SIZE); + std::cerr << "[H264Decode] Could not open input: " << errbuf << std::endl; + + // Cleanup + avio_context_free(&ioContext); + avformat_free_context(formatContext); + return (pixtype *)out; + } + + // Retrieve stream information + ret = avformat_find_stream_info(formatContext, nullptr); + if (ret < 0) + { + std::cerr << "[H264Decode] Could not find stream information" << std::endl; + + // Cleanup + avformat_close_input(&formatContext); + return (pixtype *)out; + } + + // Print some information about the media + // av_dump_format(formatContext, 0, nullptr, 0); + + // Find video stream + int video_stream_idx = -1; + for (unsigned int i = 0; i < formatContext->nb_streams; ++i) + { + if (formatContext->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) + { + video_stream_idx = i; + break; + } + } + + if (video_stream_idx == -1) + { + std::cerr << "[H264Decode] Could not find video stream" << std::endl; + return (pixtype *)out; + } + + // Get codec parameters and codec context + AVCodecParameters *codecpar = formatContext->streams[video_stream_idx]->codecpar; + AVCodec *codec = (AVCodec *)avcodec_find_decoder(codecpar->codec_id); + AVCodecContext *codec_ctx = avcodec_alloc_context3(codec); + avcodec_parameters_to_context(codec_ctx, codecpar); + + // Open codec + if (avcodec_open2(codec_ctx, codec, nullptr) < 0) + { + std::cerr << "[H264Decode] Could not open codec" << std::endl; + return (pixtype *)out; + } + + // Allocate frame and packet + AVFrame *frame = av_frame_alloc(); + AVPacket packet; + av_init_packet(&packet); + + size_t frame_cnt = 0; + + // Read frames from the video stream + while (av_read_frame(formatContext, &packet) >= 0) + { + if (packet.stream_index == video_stream_idx) + { + // Decode video frame + int ret = avcodec_send_packet(codec_ctx, &packet); + if (ret < 0) + { + std::cerr << "[H264Decode] Error sending packet for decoding" << std::endl; + break; + } + + while (ret >= 0) + { + ret = avcodec_receive_frame(codec_ctx, frame); + if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) + { + break; + } + else if (ret < 0) + { + std::cerr << "[H264Decode] Error during decoding" << std::endl; + break; + } + + for (size_t x = 0; x < sizex; x++) + { + for (size_t y = 0; y < sizey; y++) + { + const size_t in_offset = (x * frame->linesize[0]) + y; + const size_t out_offset = (x * sizey * sizez) + (y * sizez) + frame_cnt; + + out[out_offset] = frame->data[0][in_offset]; + } + } + + frame_cnt++; + } + } + } + + while (true) + { + avcodec_send_packet(codec_ctx, nullptr); + int ret = avcodec_receive_frame(codec_ctx, frame); + if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) + { + break; + } + else if (ret < 0) + { + std::cerr << "[H264Decode] Error during decoding" << std::endl; + break; + } + + for (size_t x = 0; x < sizex; x++) + { + for (size_t y = 0; y < sizey; y++) + { + const size_t in_offset = (x * frame->linesize[0]) + y; + const size_t out_offset = (x * sizey * sizez) + (y * sizez) + frame_cnt; + + out[out_offset] = frame->data[0][in_offset]; + } + } + + frame_cnt++; + } + + av_packet_unref(&packet); + + // Cleanup + avformat_close_input(&formatContext); + + return (pixtype *)out; +}