pipewirerecord.cpp

/*
    SPDX-FileCopyrightText: 2022 Aleix Pol Gonzalez <aleixpol@kde.org>
 
    SPDX-License-Identifier: LGPL-2.1-only OR LGPL-3.0-only OR LicenseRef-KDE-Accepted-LGPL
*/
 
#include "pipewirerecord.h"
#include "encoder_p.h"
#include "glhelpers.h"
#include "pipewirerecord_p.h"
#include <logging_record.h>
 
#include <QGuiApplication>
#include <QImage>
#include <QPainter>
 
#include <KShell>
 
#include <unistd.h>
extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/timestamp.h>
}
 
#undef av_err2str
 
#ifdef av_ts2str
#undef av_ts2str
char buf[AV_TS_MAX_STRING_SIZE];
#define av_ts2str(ts) av_ts_make_string(buf, ts)
#endif // av_ts2str
 
#ifdef av_ts2timestr
#undef av_ts2timestr
char timebuf[AV_TS_MAX_STRING_SIZE];
#define av_ts2timestr(ts, tb) av_ts_make_time_string(timebuf, ts, tb)
#endif // av_ts2timestr
 
static void log_packet(const AVFormatContext *fmt_ctx, const AVPacket *pkt)
{
    AVRational *time_base = &fmt_ctx->streams[pkt->stream_index]->time_base;
 
    qCDebug(PIPEWIRERECORD_LOGGING,
            "pts:%s pts_time:%s dts:%s dts_time:%s duration:%s duration_time:%s "
            "stream_index:%d",
            av_ts2str(pkt->pts),
            av_ts2timestr(pkt->pts, time_base),
            av_ts2str(pkt->dts),
            av_ts2timestr(pkt->dts, time_base),
            av_ts2str(pkt->duration),
            av_ts2timestr(pkt->duration, time_base),
            pkt->stream_index);
}
 
PipeWireRecord::PipeWireRecord(QObject *parent)
    : PipeWireBaseEncodedStream(parent)
    , d(new PipeWireRecordPrivate)
{
}
 
PipeWireRecord::~PipeWireRecord() = default;
 
void PipeWireRecord::setOutput(const QString &_output)
{
    const QString output = KShell::tildeExpand(_output);
 
    if (d->m_output == output)
        return;
 
    d->m_output = output;
    Q_EMIT outputChanged(output);
}
 
QString PipeWireRecord::output() const
{
    return d->m_output;
}
 
QString PipeWireRecord::extension() const
{
    static QHash<PipeWireBaseEncodedStream::Encoder, QString> s_extensions = {
        {PipeWireBaseEncodedStream::H264Main, QStringLiteral("mp4")},
        {PipeWireBaseEncodedStream::H264Baseline, QStringLiteral("mp4")},
        {PipeWireBaseEncodedStream::VP8, QStringLiteral("webm")},
        {PipeWireBaseEncodedStream::VP9, QStringLiteral("webm")},
        {PipeWireBaseEncodedStream::WebP, QStringLiteral("webp")},
        {PipeWireBaseEncodedStream::Gif, QStringLiteral("gif")},
    };
    return s_extensions.value(encoder());
}
 
PipeWireRecordProduce::PipeWireRecordProduce(PipeWireBaseEncodedStream::Encoder encoder, uint nodeId, uint fd, const Fraction &framerate, const QString &output)
    : PipeWireProduce(encoder, nodeId, fd, framerate)
    , m_output(output)
{
}
 
bool PipeWireRecordProduce::setupFormat()
{
    avformat_alloc_output_context2(&m_avFormatContext, nullptr, nullptr, m_output.toUtf8().constData());
    if (!m_avFormatContext) {
        qCWarning(PIPEWIRERECORD_LOGGING) << "Could not deduce output format from file: using WebM." << m_output;
        avformat_alloc_output_context2(&m_avFormatContext, nullptr, "webm", m_output.toUtf8().constData());
    }
    if (!m_avFormatContext) {
        qCDebug(PIPEWIRERECORD_LOGGING) << "could not set stream up";
        return false;
    }
 
    const Fraction framerate = m_stream->framerate();
    int ret = avio_open(&m_avFormatContext->pb, QFile::encodeName(m_output).constData(), AVIO_FLAG_WRITE);
    if (ret < 0) {
        qCWarning(PIPEWIRERECORD_LOGGING) << "Could not open" << m_output << av_err2str(ret);
        return false;
    }
 
    auto avStream = avformat_new_stream(m_avFormatContext, nullptr);
    avStream->start_time = 0;
    if (framerate) {
        avStream->r_frame_rate.num = framerate.numerator;
        avStream->r_frame_rate.den = framerate.denominator;
        avStream->avg_frame_rate.num = framerate.numerator;
        avStream->avg_frame_rate.den = framerate.denominator;
    }
 
    ret = avcodec_parameters_from_context(avStream->codecpar, m_encoder->avCodecContext());
    if (ret < 0) {
        qCWarning(PIPEWIRERECORD_LOGGING) << "Error occurred when passing the codec:" << av_err2str(ret);
        return false;
    }
 
    AVDictionary *options = nullptr;
    const auto codecId = m_avFormatContext->oformat->video_codec;
    if (codecId == AV_CODEC_ID_GIF || codecId == AV_CODEC_ID_WEBP) {
        av_dict_set_int(&options, "loop", 0, 0);
    }
    ret = avformat_write_header(m_avFormatContext, &options);
    if (ret < 0) {
        qCWarning(PIPEWIRERECORD_LOGGING) << "Error occurred when writing header:" << av_err2str(ret);
        return false;
    }
 
    return true;
}
 
void PipeWireRecordProduce::processFrame(const PipeWireFrame &frame)
{
    PipeWireProduce::processFrame(frame);
    if (frame.cursor && !frame.dmabuf && !frame.dataFrame && m_frameWithoutMetadataCursor.dataFrame) {
        m_encoder->filterFrame(m_frameWithoutMetadataCursor);
    }
}
 
void PipeWireRecordProduce::aboutToEncode(PipeWireFrame &frame)
{
    if (!frame.dataFrame) {
        return;
    }
 
    if (m_cursor.position && !m_cursor.texture.isNull()) {
        auto image = frame.dataFrame->toImage();
        // Do not copy the image if it's already ours
        if (m_frameWithoutMetadataCursor.dataFrame->cleanup != frame.dataFrame->cleanup) {
            m_frameWithoutMetadataCursor.dataFrame = frame.dataFrame->copy();
        }
        QPainter p(&image);
        p.drawImage(*m_cursor.position, m_cursor.texture);
    }
}
 
void PipeWireRecordProduce::processPacket(AVPacket *packet)
{
    packet->stream_index = (*m_avFormatContext->streams)->index;
    av_packet_rescale_ts(packet, m_encoder->avCodecContext()->time_base, (*m_avFormatContext->streams)->time_base);
    log_packet(m_avFormatContext, packet);
    auto ret = av_interleaved_write_frame(m_avFormatContext, packet);
    if (ret < 0) {
        qCWarning(PIPEWIRERECORD_LOGGING) << "Error while writing output packet:" << av_err2str(ret);
    }
}
 
std::unique_ptr<PipeWireProduce> PipeWireRecord::makeProduce()
{
    return std::make_unique<PipeWireRecordProduce>(encoder(), nodeId(), fd(), maxFramerate(), d->m_output);
}
 
int64_t PipeWireRecordProduce::framePts(const std::optional<std::chrono::nanoseconds> &presentationTimestamp)
{
    const auto current = std::chrono::duration_cast<std::chrono::milliseconds>(*presentationTimestamp).count();
    if ((*m_avFormatContext->streams)->start_time == 0) {
        (*m_avFormatContext->streams)->start_time = current;
    }
 
    return current - (*m_avFormatContext->streams)->start_time;
}
 
void PipeWireRecordProduce::cleanup()
{
    if (m_avFormatContext) {
        if (auto result = av_write_trailer(m_avFormatContext); result < 0) {
            qCWarning(PIPEWIRERECORD_LOGGING) << "Could not write trailer";
        }
 
        avio_closep(&m_avFormatContext->pb);
        avformat_free_context(m_avFormatContext);
    }
}
 
#include "moc_pipewirerecord.cpp"
 
#include "moc_pipewirerecord_p.cpp"