screen-recorder/obs-manager.cpp

#include "obs-manager.hpp"
#include "settings-manager.hpp"

#include <ctime>
#include <list>
#include <filesystem>
#include <stdexcept>

using namespace std;

static void obs_render(void *param, uint32_t cx, uint32_t cy)
{
    obs_render_main_texture();
}

static void OBSStartRecording(void *data, calldata_t *params)
{
    OBSManager *o = static_cast<OBSManager*>(data);
    o->startRecording.emit();
}

static void OBSStopRecording(void *data, calldata_t *params)
{
    OBSManager *o = static_cast<OBSManager*>(data);
    o->stopRecording.emit();
}

OBSManager::OBSManager()
{
    mPlugins = {
        "obs-ffmpeg.so",
        "obs-outputs.so",
        "obs-x264.so",
        "linux-v4l2.so",
        "linux-capture.so",
    };
}

OBSManager::~OBSManager()
{
    Cleanup();
}

string OBSManager::GetVersion()
{
    return string(obs_get_version_string());
}

void OBSManager::Initialize()
{
    if (isInitialized)
        return;
    
    mSources = list<OBSSource>();
    auto settings = new SettingsManager();
    LoadSettings(settings);
    
    if (!obs_startup("en-US", NULL, NULL))
        throw runtime_error("Failed to initialize OBS");
    
    loadPlugins();
    
    obs_video_info v = {};
    v.graphics_module = "libobs-opengl.so.0";
    v.fps_num = 30000;
    v.fps_den = 1001;
    v.base_width = PreviewWidth;
    v.base_height = PreviewHeight;
    v.output_width = PreviewWidth;
    v.output_height = PreviewHeight;
    v.output_format = VIDEO_FORMAT_NV12;
    v.adapter = 0;
    v.gpu_conversion = true;
    v.colorspace = VIDEO_CS_601;
    v.range = VIDEO_RANGE_PARTIAL;
    v.scale_type = OBS_SCALE_BICUBIC;

    obs_audio_info a = {};
    a.samples_per_sec = 44100;
    a.speakers = SPEAKERS_STEREO;

    obs_reset_video(&v);
    obs_reset_audio(&a);
    
    isInitialized = true;
}

void OBSManager::SetPreviewWindow(XID wid, Display *wdisplay)
{
    auto settings = new SettingsManager();
    
    gs_init_data init = {};
    init.cx = PreviewWidth;
    init.cy = PreviewHeight;
    init.format = GS_BGRA;
    init.zsformat = GS_ZS_NONE;
    init.window.id = wid;
    init.window.display = (void*)wdisplay;
    
    mDisplay = obs_display_create(&init, 0);
    if (mDisplay == nullptr)
        throw runtime_error("Failed to create display");
    
    obs_display_add_draw_callback(mDisplay, obs_render, nullptr);
    
    OBSScene scene = obs_scene_create("scene1");
    if (scene == NULL)
        throw runtime_error("Couldn't create scene\n");
    
    if (settings->GetBool(SETTINGS_KEY_SCREEN_ENABLED))
    {
        auto source = CreateScreenSource();
        obs_scene_add(scene, source);
        mSources.push_back(source);
    }
    
    if (settings->GetBool(SETTINGS_KEY_WEBCAM_ENABLED))
    {
        vec2 scale;
        vec2_set(&scale, 0.5f, 0.5f);
        
        auto source = CreateWebcamSource();
        auto item = obs_scene_add(scene, source);
        obs_sceneitem_set_scale(item, &scale);
        mSources.push_back(source);
    }
    
    if (settings->GetBool(SETTINGS_KEY_AUDIO_ENABLED))
    {
        auto source = CreateAudioSource();
        obs_scene_add(scene, source);
        mSources.push_back(source);
    }
    
    obs_set_output_source(0, obs_scene_get_source(scene));
}

void OBSManager::StartRecording()
{
    if (isRecording)
        return;
    
    auto s = new SettingsManager();
    
    obs_encoder_t *venc = obs_video_encoder_create("obs_x264", "x264", nullptr, nullptr);
    obs_encoder_t *aenc = obs_audio_encoder_create("ffmpeg_aac", "aac", nullptr, 0, nullptr);
    
    obs_encoder_set_video(venc, obs_get_video());
    obs_encoder_set_audio(aenc, obs_get_audio());
    
    string path = s->GetWithDefault(SETTINGS_KEY_OUTPUT_DIR, filesystem::current_path());
    if (path.back() != '/')
        path += "/";
    
    string fileName = path + "recording_" + to_string(time(0)) + ".mp4";
    
    obs_data_t *settings = obs_data_create();
    obs_data_set_string(settings, "directory", path.c_str());
    obs_data_set_string(settings, "url", fileName.c_str());
    
    mOutput = obs_output_create("ffmpeg_output", "ffmpeg_output", nullptr, nullptr);
    obs_output_set_video_encoder(mOutput, venc);
    obs_output_set_audio_encoder(mOutput, aenc, 0);
    obs_output_update(mOutput, settings);
    obs_output_set_media(mOutput, obs_get_video(), obs_get_audio());
    
    obs_data_release(settings);
    if (!obs_output_start(mOutput))
        throw runtime_error("Failed to start recording");
    
    obsStartRecording.Connect(obs_output_get_signal_handler(mOutput), "start", OBSStartRecording, this);
    obsStopRecording.Connect(obs_output_get_signal_handler(mOutput), "stop", OBSStopRecording, this);
    
    isRecording = true;
}

void OBSManager::StopRecording()
{
    if (isRecording && mOutput != nullptr)
        obs_output_stop(mOutput);
    
    isRecording = false;
}

OBSSource OBSManager::CreateScreenSource()
{
    obs_data_t *settings = obs_data_create();
    OBSSource source = obs_source_create("xshm_input", "Screen Source", settings, NULL);
    if (source == NULL)
        throw runtime_error("Couldn't create screen source");
    
    obs_data_release(settings);
    return source;
}

OBSSource OBSManager::CreateWebcamSource()
{
    obs_data_t *settings = obs_data_create();
    obs_data_set_string(settings, "device_id", mWebcamDeviceID.c_str());

    OBSSource source = obs_source_create("v4l2_input", "Webcam Source", settings, NULL);
    if (source == NULL)
        throw runtime_error("Couldn't create webcam source");
    
    obs_data_release(settings);
    return source;
}

OBSSource OBSManager::CreateAudioSource()
{
    obs_data_t *settings = obs_data_create();
    OBSSource source = obs_source_create("audio_line", "Audio Source", settings, NULL);
    if (source == NULL)
        throw runtime_error("Couldn't create screen source");
    
    obs_data_release(settings);
    return source;
}

void OBSManager::Cleanup()
{
    if (!isInitialized)
        return;
    
    StopRecording();
    
    for (auto source : mSources)
    {
        obs_source_remove(source);
    }
    
    if (mDisplay != nullptr)
        obs_display_destroy(mDisplay);
    
    obs_shutdown();
    isInitialized = false;
}

void OBSManager::LoadSettings(SettingsManager *settings)
{
    mPluginDir = settings->GetWithDefault(SETTINGS_KEY_PLUGIN_DIR, SETTINGS_DEFAULT_PLUGIN_DIR);
    mOutputDir = settings->GetWithDefault(SETTINGS_KEY_OUTPUT_DIR, std::filesystem::current_path());
    mWebcamDeviceID = settings->Get(SETTINGS_KEY_VIDEO_DEVICE_ID);
}

bool OBSManager::IsRecording()
{
    return isRecording;
}

void OBSManager::loadPlugin(string name)
{
    obs_module_t *module;
    
    string path;
    path.append(mPluginDir);
    path.append(name);
    
    int res = obs_open_module(&module, path.c_str(), nullptr);
    if (res != MODULE_SUCCESS)
        throw runtime_error("Failed to open plugin");
    
    obs_init_module(module);
}


void OBSManager::loadPlugins()
{
    for (string plugin : mPlugins)
    {
        loadPlugin(plugin);
    }
    
    obs_post_load_modules();
}

void OBSManager::printTypes()
{
    const char *t;
    for (int i = 0; obs_enum_input_types(i, &t); ++i)
    {
        cout << "input type: " << t << endl;
    }
    
    for (int i = 0; obs_enum_output_types(i, &t); ++i)
    {
        cout << "output type: " << t << endl;
    }
    
    for (int i = 0; obs_enum_encoder_types(i, &t); ++i)
    {
        cout << "encoder type: " << t << endl;
    }
}