Skip to content

Commit

Permalink
refactor(*): refactored the encoder and decoder into classes
Browse files Browse the repository at this point in the history 
Signed-off-by: k4yt3x <[email protected]>
  • Loading branch information
@k4yt3x
k4yt3x committed Nov 17, 2024
1 parent b520d51 commit 169509b
Show file tree
Hide file tree
Showing 6 changed files with 321 additions and 350 deletions.
24 changes: 14 additions & 10 deletions CMakeLists.txt
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,10 @@
cmake_minimum_required(VERSION 3.10)
project(video2x VERSION 6.1.1 LANGUAGES CXX)

if(POLICY CMP0167)
cmake_policy(SET CMP0167 NEW)
endif()

# Set the C++ standard
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
Expand All @@ -13,19 +17,19 @@ endif()

# Set the default optimization flags for Release builds
if(CMAKE_BUILD_TYPE STREQUAL "Release")
if (CMAKE_CXX_COMPILER_ID STREQUAL "MSVC")
if(CMAKE_CXX_COMPILER_ID STREQUAL "MSVC")
set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} /Ox /GL /LTCG /MD /DNDEBUG")
elseif (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR CMAKE_CXX_COMPILER_ID MATCHES "Clang")
elseif(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR CMAKE_CXX_COMPILER_ID MATCHES "Clang")
set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -O3 -march=native -flto")
set(CMAKE_EXE_LINKER_FLAGS_RELEASE "${CMAKE_EXE_LINKER_FLAGS_RELEASE} -s")
set(CMAKE_SHARED_LINKER_FLAGS_RELEASE "${CMAKE_SHARED_LINKER_FLAGS_RELEASE} -s")
endif()
endif()

# Set global compile options for all targets
if (CMAKE_CXX_COMPILER_ID STREQUAL "MSVC")
if(CMAKE_CXX_COMPILER_ID STREQUAL "MSVC")
add_compile_options(/W4 /permissive-)
elseif (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR CMAKE_CXX_COMPILER_ID MATCHES "Clang")
elseif(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR CMAKE_CXX_COMPILER_ID MATCHES "Clang")
add_compile_options(-Wall -Wextra -Wpedantic -Wconversion -Wshadow)
endif()

Expand Down Expand Up @@ -70,7 +74,7 @@ if(WIN32)
set(SPIRV_BUILD_PATH
${CMAKE_BINARY_DIR}/realesrgan-prefix/src/realesrgan-build/ncnn/glslang/SPIRV
)
if (CMAKE_BUILD_TYPE STREQUAL "Release")
if(CMAKE_BUILD_TYPE STREQUAL "Release")
set(SPIRV_LIB ${SPIRV_BUILD_PATH}/Release/SPIRV.lib)
else()
set(SPIRV_LIB ${SPIRV_BUILD_PATH}/Debug/SPIRVd.lib)
Expand Down Expand Up @@ -203,7 +207,7 @@ else()
endif()

# spdlog
if (USE_SYSTEM_SPDLOG)
if(USE_SYSTEM_SPDLOG)
find_package(spdlog REQUIRED)
list(APPEND ALL_INCLUDE_DIRS ${spdlog_INCLUDE_DIRS})
set(SPDLOG_LIB spdlog::spdlog)
Expand All @@ -214,7 +218,7 @@ endif()
list(APPEND ALL_LIBRARIES ${SPDLOG_LIB})

# Boost
if (USE_SYSTEM_BOOST)
if(USE_SYSTEM_BOOST)
find_package(Boost REQUIRED COMPONENTS program_options)
list(APPEND ALL_INCLUDE_DIRS ${Boost_INCLUDE_DIRS})
else()
Expand All @@ -228,7 +232,7 @@ else()
endif()
set(BOOST_LIB Boost::program_options)

if (BUILD_VIDEO2X_CLI)
if(BUILD_VIDEO2X_CLI)
find_package(Vulkan REQUIRED)
set(VULKAN_LIB Vulkan::Vulkan)
endif()
Expand Down Expand Up @@ -295,15 +299,15 @@ endif()
target_link_libraries(libvideo2x PRIVATE ${ALL_LIBRARIES})

if(NOT WIN32)
if (USE_SYSTEM_NCNN)
if(USE_SYSTEM_NCNN)
target_link_libraries(libvideo2x PUBLIC ncnn)
else()
target_link_libraries(libvideo2x PRIVATE ncnn)
endif()
endif()

# Create the executable 'video2x'
if (BUILD_VIDEO2X_CLI)
if(BUILD_VIDEO2X_CLI)
file(GLOB VIDEO2X_SOURCES tools/video2x/src/*.cpp)
add_executable(video2x ${VIDEO2X_SOURCES})
set_target_properties(video2x PROPERTIES OUTPUT_NAME video2x)
Expand Down
28 changes: 20 additions & 8 deletions include/libvideo2x/decoder.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,13 +8,25 @@ extern "C" {
#include <libavformat/avformat.h>
}

int init_decoder(
AVHWDeviceType hw_type,
AVBufferRef *hw_ctx,
std::filesystem::path in_fpath,
AVFormatContext **fmt_ctx,
AVCodecContext **dec_ctx,
int *in_vstream_idx
);
class Decoder {
public:
Decoder();
~Decoder();

int init(AVHWDeviceType hw_type, AVBufferRef *hw_ctx, const std::filesystem::path &in_fpath);

AVFormatContext *get_format_context() const;
AVCodecContext *get_codec_context() const;
int get_video_stream_index() const;

private:
static enum AVPixelFormat hw_pix_fmt_;
static enum AVPixelFormat
get_hw_format(AVCodecContext *ctx, const enum AVPixelFormat *pix_fmts);

AVFormatContext *fmt_ctx_;
AVCodecContext *dec_ctx_;
int in_vstream_idx_;
};

#endif // DECODER_H
58 changes: 32 additions & 26 deletions include/libvideo2x/encoder.h
View file
Original file line number Diff line number Diff line change
@@ -1,37 +1,43 @@
#ifndef ENCODER_H
#define ENCODER_H

#include <cstdint>
#include <filesystem>

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/opt.h>
#include <libavutil/pixdesc.h>
}

#include "libvideo2x.h"

int init_encoder(
AVBufferRef *hw_ctx,
std::filesystem::path out_fpath,
AVFormatContext *ifmt_ctx,
AVFormatContext **ofmt_ctx,
AVCodecContext **enc_ctx,
AVCodecContext *dec_ctx,
EncoderConfig *encoder_config,
int in_vstream_idx,
int *out_vstream_idx,
int **stream_map
);

int write_frame(
AVFrame *frame,
AVCodecContext *enc_ctx,
AVFormatContext *ofmt_ctx,
int out_vstream_idx,
int64_t frame_idx
);

int flush_encoder(AVCodecContext *enc_ctx, AVFormatContext *ofmt_ctx, int out_vstream_idx);
#include "libvideo2x/libvideo2x.h"

class Encoder {
public:
Encoder();
~Encoder();

int init(
AVBufferRef *hw_ctx,
const std::filesystem::path &out_fpath,
AVFormatContext *ifmt_ctx,
AVCodecContext *dec_ctx,
EncoderConfig *encoder_config,
int in_vstream_idx
);

int write_frame(AVFrame *frame, int64_t frame_idx);
int flush();

AVCodecContext *get_encoder_context() const;
AVFormatContext *get_format_context() const;
int *get_stream_map() const;
int get_output_video_stream_index() const;

private:
AVFormatContext *ofmt_ctx_;
AVCodecContext *enc_ctx_;
int out_vstream_idx_;
int *stream_map_;
};

#endif // ENCODER_H
100 changes: 58 additions & 42 deletions src/decoder.cpp
View file
Original file line number Diff line number Diff line change
@@ -1,57 +1,62 @@
#include "decoder.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <spdlog/spdlog.h>

static enum AVPixelFormat hw_pix_fmt = AV_PIX_FMT_NONE;
enum AVPixelFormat Decoder::hw_pix_fmt_ = AV_PIX_FMT_NONE;

Decoder::Decoder() : fmt_ctx_(nullptr), dec_ctx_(nullptr), in_vstream_idx_(-1) {}

Decoder::~Decoder() {
if (dec_ctx_) {
avcodec_free_context(&dec_ctx_);
dec_ctx_ = nullptr;
}
if (fmt_ctx_) {
avformat_close_input(&fmt_ctx_);
fmt_ctx_ = nullptr;
}
}

// Callback function to choose the hardware-accelerated pixel format
static enum AVPixelFormat get_hw_format(AVCodecContext *_, const enum AVPixelFormat *pix_fmts) {
enum AVPixelFormat Decoder::get_hw_format(AVCodecContext *_, const enum AVPixelFormat *pix_fmts) {
for (const enum AVPixelFormat *p = pix_fmts; *p != AV_PIX_FMT_NONE; p++) {
if (*p == hw_pix_fmt) {
if (*p == hw_pix_fmt_) {
return *p;
}
}
spdlog::error("Failed to get HW surface format.");
return AV_PIX_FMT_NONE;
}

int init_decoder(
int Decoder::init(
AVHWDeviceType hw_type,
AVBufferRef *hw_ctx,
std::filesystem::path in_fpath,
AVFormatContext **fmt_ctx,
AVCodecContext **dec_ctx,
int *in_vstream_idx
const std::filesystem::path &in_fpath
) {
AVFormatContext *ifmt_ctx = NULL;
AVCodecContext *codec_ctx = NULL;
int ret;

if ((ret = avformat_open_input(&ifmt_ctx, in_fpath.u8string().c_str(), NULL, NULL)) < 0) {
spdlog::error("Could not open input file '{}'", in_fpath.u8string().c_str());
// Open the input file
if ((ret = avformat_open_input(&fmt_ctx_, in_fpath.u8string().c_str(), nullptr, nullptr)) < 0) {
spdlog::error("Could not open input file '{}'", in_fpath.u8string());
return ret;
}

if ((ret = avformat_find_stream_info(ifmt_ctx, NULL)) < 0) {
// Retrieve stream information
if ((ret = avformat_find_stream_info(fmt_ctx_, nullptr)) < 0) {
spdlog::error("Failed to retrieve input stream information");
return ret;
}

// Find the first video stream
ret = av_find_best_stream(ifmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, NULL, 0);
ret = av_find_best_stream(fmt_ctx_, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
if (ret < 0) {
spdlog::error("Could not find video stream in the input file");
return ret;
}

int stream_index = ret;
AVStream *video_stream = ifmt_ctx->streams[stream_index];
AVStream *video_stream = fmt_ctx_->streams[stream_index];

// Set up the decoder
// Find the decoder for the video stream
const AVCodec *decoder = avcodec_find_decoder(video_stream->codecpar->codec_id);
if (!decoder) {
spdlog::error(
Expand All @@ -61,16 +66,28 @@ int init_decoder(
return AVERROR_DECODER_NOT_FOUND;
}

codec_ctx = avcodec_alloc_context3(decoder);
if (!codec_ctx) {
// Allocate the decoder context
dec_ctx_ = avcodec_alloc_context3(decoder);
if (!dec_ctx_) {
spdlog::error("Failed to allocate the decoder context");
return AVERROR(ENOMEM);
}

// Copy codec parameters from input stream to decoder context
if ((ret = avcodec_parameters_to_context(dec_ctx_, video_stream->codecpar)) < 0) {
spdlog::error("Failed to copy decoder parameters to input decoder context");
return ret;
}

// Set the time base and frame rate
dec_ctx_->time_base = video_stream->time_base;
dec_ctx_->pkt_timebase = video_stream->time_base;
dec_ctx_->framerate = av_guess_frame_rate(fmt_ctx_, video_stream, nullptr);

// Set hardware device context
if (hw_ctx != nullptr) {
codec_ctx->hw_device_ctx = av_buffer_ref(hw_ctx);
codec_ctx->get_format = get_hw_format;
dec_ctx_->hw_device_ctx = av_buffer_ref(hw_ctx);
dec_ctx_->get_format = get_hw_format;

// Automatically determine the hardware pixel format
for (int i = 0;; i++) {
Expand All @@ -81,36 +98,35 @@ int init_decoder(
decoder->name,
av_hwdevice_get_type_name(hw_type)
);
avcodec_free_context(&codec_ctx);
avformat_close_input(&ifmt_ctx);
return AVERROR(ENOSYS);
}
if (config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX &&
config->device_type == hw_type) {
hw_pix_fmt = config->pix_fmt;
hw_pix_fmt_ = config->pix_fmt;
break;
}
}
}

if ((ret = avcodec_parameters_to_context(codec_ctx, video_stream->codecpar)) < 0) {
spdlog::error("Failed to copy decoder parameters to input decoder context");
return ret;
}

// Set decoder time base and frame rate
codec_ctx->time_base = video_stream->time_base;
codec_ctx->pkt_timebase = video_stream->time_base;
codec_ctx->framerate = av_guess_frame_rate(ifmt_ctx, video_stream, NULL);

if ((ret = avcodec_open2(codec_ctx, decoder, NULL)) < 0) {
// Open the decoder
if ((ret = avcodec_open2(dec_ctx_, decoder, nullptr)) < 0) {
spdlog::error("Failed to open decoder for stream #{}", stream_index);
return ret;
}

*fmt_ctx = ifmt_ctx;
*dec_ctx = codec_ctx;
*in_vstream_idx = stream_index;
in_vstream_idx_ = stream_index;

return 0;
}

AVFormatContext *Decoder::get_format_context() const {
return fmt_ctx_;
}

AVCodecContext *Decoder::get_codec_context() const {
return dec_ctx_;
}

int Decoder::get_video_stream_index() const {
return in_vstream_idx_;
}
Loading

0 comments on commit 169509b

Please sign in to comment.