feat: translate subtitles with an OpenAI-compatible LLM
This commit is contained in:
tri committed
1 parent
48cac86d60
commit
4665bf632a
21 files changed
+1877
-5
No files matched your search
@@ -0,0 +1,141 @@
|
||||
#include "subtitleextractor.h"
|
||||
|
||||
#include <KLocalizedString>
|
||||
|
||||
#include <memory>
|
||||
|
||||
extern "C" {
|
||||
#include <libavcodec/avcodec.h>
|
||||
#include <libavformat/avformat.h>
|
||||
#include <libavutil/error.h>
|
||||
}
|
||||
|
||||
namespace
|
||||
{
|
||||
struct FormatContextDeleter {
|
||||
void operator()(AVFormatContext *context) const
|
||||
{
|
||||
avformat_close_input(&context);
|
||||
}
|
||||
};
|
||||
struct CodecContextDeleter {
|
||||
void operator()(AVCodecContext *context) const
|
||||
{
|
||||
avcodec_free_context(&context);
|
||||
}
|
||||
};
|
||||
struct PacketDeleter {
|
||||
void operator()(AVPacket *packet) const
|
||||
{
|
||||
av_packet_free(&packet);
|
||||
}
|
||||
};
|
||||
|
||||
QString ffmpegError(int code)
|
||||
{
|
||||
char buffer[AV_ERROR_MAX_STRING_SIZE] = {};
|
||||
av_strerror(code, buffer, sizeof(buffer));
|
||||
return QString::fromUtf8(buffer);
|
||||
}
|
||||
|
||||
qint64 toMs(int64_t value, AVRational timeBase)
|
||||
{
|
||||
return av_rescale_q(value, timeBase, AVRational{1, 1000});
|
||||
}
|
||||
|
||||
// ASS events from the decoder look like "ReadOrder,Layer,Style,Name,MarginL,MarginR,MarginV,Effect,Text".
|
||||
QString assRectText(const char *ass)
|
||||
{
|
||||
const QString event = QString::fromUtf8(ass);
|
||||
qsizetype from = 0;
|
||||
for (int field = 0; field < 8; ++field) {
|
||||
from = event.indexOf(u',', from);
|
||||
if (from < 0) {
|
||||
return {};
|
||||
}
|
||||
++from;
|
||||
}
|
||||
return Subtitles::assEventText(QStringView(event).mid(from));
|
||||
}
|
||||
}
|
||||
|
||||
namespace Subtitles
|
||||
{
|
||||
ExtractResult extractEmbedded(const QString &path, int streamIndex, const std::atomic_bool &cancelled)
|
||||
{
|
||||
AVFormatContext *rawFormat = nullptr;
|
||||
if (const int ret = avformat_open_input(&rawFormat, path.toUtf8().constData(), nullptr, nullptr); ret < 0) {
|
||||
return {{}, i18n("Could not open %1: %2", path, ffmpegError(ret))};
|
||||
}
|
||||
std::unique_ptr<AVFormatContext, FormatContextDeleter> format(rawFormat);
|
||||
|
||||
if (streamIndex < 0 || streamIndex >= static_cast<int>(format->nb_streams)) {
|
||||
return {{}, i18n("Subtitle stream %1 does not exist in this file.", streamIndex)};
|
||||
}
|
||||
AVStream *stream = format->streams[streamIndex];
|
||||
|
||||
const AVCodecDescriptor *descriptor = avcodec_descriptor_get(stream->codecpar->codec_id);
|
||||
if (stream->codecpar->codec_type != AVMEDIA_TYPE_SUBTITLE || !descriptor) {
|
||||
return {{}, i18n("Stream %1 is not a subtitle stream.", streamIndex)};
|
||||
}
|
||||
if (descriptor->props & AV_CODEC_PROP_BITMAP_SUB) {
|
||||
return {{}, i18n("This subtitle track is made of images (%1) and can't be translated.", QString::fromUtf8(descriptor->name))};
|
||||
}
|
||||
|
||||
const AVCodec *codec = avcodec_find_decoder(stream->codecpar->codec_id);
|
||||
if (!codec) {
|
||||
return {{}, i18n("No decoder for subtitle format %1.", QString::fromUtf8(descriptor->name))};
|
||||
}
|
||||
std::unique_ptr<AVCodecContext, CodecContextDeleter> decoder(avcodec_alloc_context3(codec));
|
||||
avcodec_parameters_to_context(decoder.get(), stream->codecpar);
|
||||
decoder->pkt_timebase = stream->time_base;
|
||||
if (const int ret = avcodec_open2(decoder.get(), codec, nullptr); ret < 0) {
|
||||
return {{}, i18n("Could not open the subtitle decoder: %1", ffmpegError(ret))};
|
||||
}
|
||||
|
||||
// Only the subtitle stream matters; let the demuxer skip the rest.
|
||||
for (unsigned i = 0; i < format->nb_streams; ++i) {
|
||||
if (static_cast<int>(i) != streamIndex) {
|
||||
format->streams[i]->discard = AVDISCARD_ALL;
|
||||
}
|
||||
}
|
||||
|
||||
ExtractResult result;
|
||||
std::unique_ptr<AVPacket, PacketDeleter> packet(av_packet_alloc());
|
||||
while (av_read_frame(format.get(), packet.get()) >= 0) {
|
||||
if (cancelled) {
|
||||
return {};
|
||||
}
|
||||
if (packet->stream_index != streamIndex || packet->pts == AV_NOPTS_VALUE) {
|
||||
av_packet_unref(packet.get());
|
||||
continue;
|
||||
}
|
||||
|
||||
AVSubtitle subtitle{};
|
||||
int gotSubtitle = 0;
|
||||
if (avcodec_decode_subtitle2(decoder.get(), &subtitle, &gotSubtitle, packet.get()) >= 0 && gotSubtitle) {
|
||||
const qint64 start = toMs(packet->pts, stream->time_base);
|
||||
qint64 end = start + toMs(packet->duration, stream->time_base);
|
||||
if (packet->duration <= 0) {
|
||||
end = subtitle.end_display_time > 0 && subtitle.end_display_time != UINT32_MAX ? start + subtitle.end_display_time : start + 3000;
|
||||
}
|
||||
|
||||
for (unsigned i = 0; i < subtitle.num_rects; ++i) {
|
||||
const AVSubtitleRect *rect = subtitle.rects[i];
|
||||
QString text;
|
||||
if (rect->type == SUBTITLE_ASS && rect->ass) {
|
||||
text = assRectText(rect->ass);
|
||||
} else if (rect->type == SUBTITLE_TEXT && rect->text) {
|
||||
text = QString::fromUtf8(rect->text).trimmed();
|
||||
}
|
||||
if (!text.isEmpty()) {
|
||||
result.cues.append({start, end, text});
|
||||
}
|
||||
}
|
||||
avsubtitle_free(&subtitle);
|
||||
}
|
||||
av_packet_unref(packet.get());
|
||||
}
|
||||
return result;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
#pragma once
|
||||
|
||||
#include "subtitleparser.h"
|
||||
|
||||
#include <atomic>
|
||||
|
||||
namespace Subtitles
|
||||
{
|
||||
struct ExtractResult {
|
||||
QList<SubtitleCue> cues;
|
||||
QString error;
|
||||
};
|
||||
|
||||
// Decodes a text subtitle stream of a media file with FFmpeg. \a streamIndex is the
|
||||
// FFmpeg stream index (mpv reports it as "ff-index"). Blocking; run it off the GUI thread.
|
||||
ExtractResult extractEmbedded(const QString &path, int streamIndex, const std::atomic_bool &cancelled);
|
||||
}
|
||||
@@ -0,0 +1,242 @@
|
||||
#include "subtitleparser.h"
|
||||
|
||||
#include <KLocalizedString>
|
||||
|
||||
#include <QFile>
|
||||
#include <QHash>
|
||||
#include <QRegularExpression>
|
||||
|
||||
#include <algorithm>
|
||||
#include <optional>
|
||||
|
||||
using namespace Qt::Literals::StringLiterals;
|
||||
|
||||
namespace
|
||||
{
|
||||
// Gap allowed between two pieces of the same line when merging them.
|
||||
constexpr qint64 MergeGapMs = 50;
|
||||
// Nobody can read a line shown for less than this; such cues are animation frames.
|
||||
constexpr qint64 MinReadableMs = 250;
|
||||
|
||||
// "01:02:03,456", "01:02:03.456", "02:03.456" (WebVTT may omit hours).
|
||||
std::optional<qint64> parseSrtTime(QStringView text)
|
||||
{
|
||||
static const QRegularExpression re(u"^(?:(\\d+):)?(\\d{1,2}):(\\d{1,2})[,.](\\d{1,3})$"_s);
|
||||
const auto match = re.matchView(text.trimmed());
|
||||
if (!match.hasMatch()) {
|
||||
return std::nullopt;
|
||||
}
|
||||
QString fraction = match.captured(4);
|
||||
while (fraction.size() < 3) {
|
||||
fraction.append(u'0');
|
||||
}
|
||||
return ((match.captured(1).toLongLong() * 60 + match.captured(2).toLongLong()) * 60 + match.captured(3).toLongLong()) * 1000 + fraction.toLongLong();
|
||||
}
|
||||
|
||||
// "1:02:03.45" (centiseconds).
|
||||
std::optional<qint64> parseAssTime(QStringView text)
|
||||
{
|
||||
static const QRegularExpression re(u"^(\\d+):(\\d{1,2}):(\\d{1,2})\\.(\\d{1,3})$"_s);
|
||||
const auto match = re.matchView(text.trimmed());
|
||||
if (!match.hasMatch()) {
|
||||
return std::nullopt;
|
||||
}
|
||||
QString fraction = match.captured(4);
|
||||
while (fraction.size() < 3) {
|
||||
fraction.append(u'0');
|
||||
}
|
||||
return ((match.captured(1).toLongLong() * 60 + match.captured(2).toLongLong()) * 60 + match.captured(3).toLongLong()) * 1000 + fraction.toLongLong();
|
||||
}
|
||||
|
||||
QString formatSrtTime(qint64 ms)
|
||||
{
|
||||
ms = std::max<qint64>(ms, 0);
|
||||
return u"%1:%2:%3,%4"_s.arg(ms / 3600000, 2, 10, u'0')
|
||||
.arg(ms / 60000 % 60, 2, 10, u'0')
|
||||
.arg(ms / 1000 % 60, 2, 10, u'0')
|
||||
.arg(ms % 1000, 3, 10, u'0');
|
||||
}
|
||||
|
||||
// Trims every line and drops empty ones.
|
||||
QString normalizeLines(const QString &text)
|
||||
{
|
||||
QStringList lines;
|
||||
for (const QString &line : text.split(u'\n')) {
|
||||
const QString trimmed = line.simplified();
|
||||
if (!trimmed.isEmpty()) {
|
||||
lines.append(trimmed);
|
||||
}
|
||||
}
|
||||
return lines.join(u'\n');
|
||||
}
|
||||
|
||||
QString stripSrtFormatting(const QString &text)
|
||||
{
|
||||
static const QRegularExpression htmlTag(u"<[^>]*>"_s);
|
||||
static const QRegularExpression assOverride(u"\\{\\\\[^}]*\\}"_s);
|
||||
QString result = text;
|
||||
result.remove(htmlTag);
|
||||
result.remove(assOverride);
|
||||
return normalizeLines(result);
|
||||
}
|
||||
}
|
||||
|
||||
namespace Subtitles
|
||||
{
|
||||
QList<SubtitleCue> parseSrt(const QString &content)
|
||||
{
|
||||
QList<SubtitleCue> cues;
|
||||
const QStringList lines = QString(content).remove(u'\r').split(u'\n');
|
||||
|
||||
for (qsizetype i = 0; i < lines.size(); ++i) {
|
||||
const QString &line = lines.at(i);
|
||||
const qsizetype arrow = line.indexOf(u"-->"_s);
|
||||
if (arrow < 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// WebVTT puts cue settings after the end time: "00:01.000 --> 00:02.000 align:start".
|
||||
const QString endPart = line.mid(arrow + 3).trimmed().section(u' ', 0, 0);
|
||||
const auto start = parseSrtTime(QStringView(line).left(arrow));
|
||||
const auto end = parseSrtTime(endPart);
|
||||
if (!start || !end) {
|
||||
continue;
|
||||
}
|
||||
|
||||
QStringList textLines;
|
||||
while (i + 1 < lines.size() && !lines.at(i + 1).trimmed().isEmpty()) {
|
||||
textLines.append(lines.at(++i));
|
||||
}
|
||||
cues.append({*start, *end, stripSrtFormatting(textLines.join(u'\n'))});
|
||||
}
|
||||
return cues;
|
||||
}
|
||||
|
||||
QList<SubtitleCue> parseAss(const QString &content)
|
||||
{
|
||||
QList<SubtitleCue> cues;
|
||||
bool inEvents = false;
|
||||
QStringList format;
|
||||
|
||||
for (const QString &rawLine : QString(content).remove(u'\r').split(u'\n')) {
|
||||
const QString line = rawLine.trimmed();
|
||||
if (line.startsWith(u'[')) {
|
||||
inEvents = line.compare(u"[Events]"_s, Qt::CaseInsensitive) == 0;
|
||||
continue;
|
||||
}
|
||||
if (!inEvents) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (line.startsWith(u"Format:"_s)) {
|
||||
format.clear();
|
||||
for (const QString &field : line.mid(7).split(u',')) {
|
||||
format.append(field.trimmed());
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if (!line.startsWith(u"Dialogue:"_s) || format.isEmpty()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Text is the last field and may itself contain commas.
|
||||
const QString values = line.mid(9).trimmed();
|
||||
QStringList fields;
|
||||
qsizetype from = 0;
|
||||
for (qsizetype n = 0; n < format.size() - 1; ++n) {
|
||||
const qsizetype comma = values.indexOf(u',', from);
|
||||
if (comma < 0) {
|
||||
break;
|
||||
}
|
||||
fields.append(values.mid(from, comma - from).trimmed());
|
||||
from = comma + 1;
|
||||
}
|
||||
if (fields.size() != format.size() - 1) {
|
||||
continue;
|
||||
}
|
||||
fields.append(values.mid(from));
|
||||
|
||||
const auto start = parseAssTime(fields.value(format.indexOf(u"Start"_s)));
|
||||
const auto end = parseAssTime(fields.value(format.indexOf(u"End"_s)));
|
||||
if (!start || !end) {
|
||||
continue;
|
||||
}
|
||||
cues.append({*start, *end, assEventText(fields.last())});
|
||||
}
|
||||
return cues;
|
||||
}
|
||||
|
||||
QList<SubtitleCue> parseFile(const QString &path, QString *error)
|
||||
{
|
||||
QFile file(path);
|
||||
if (!file.open(QIODevice::ReadOnly)) {
|
||||
if (error) {
|
||||
*error = i18n("Could not read %1: %2", path, file.errorString());
|
||||
}
|
||||
return {};
|
||||
}
|
||||
|
||||
QString content = QString::fromUtf8(file.readAll());
|
||||
if (content.startsWith(QChar(0xFEFF))) {
|
||||
content.remove(0, 1);
|
||||
}
|
||||
if (content.contains(u"[Events]"_s, Qt::CaseInsensitive)) {
|
||||
return parseAss(content);
|
||||
}
|
||||
return parseSrt(content);
|
||||
}
|
||||
|
||||
QString assEventText(QStringView text)
|
||||
{
|
||||
static const QRegularExpression overrideBlock(u"\\{[^}]*\\}"_s);
|
||||
// \p1, \p2… switch to drawing mode: the "text" is vector shapes, not words.
|
||||
static const QRegularExpression drawing(u"\\{[^}]*\\\\p[1-9][^}]*\\}"_s);
|
||||
|
||||
QString result = text.toString();
|
||||
if (drawing.match(result).hasMatch()) {
|
||||
return {};
|
||||
}
|
||||
result.remove(overrideBlock);
|
||||
result.replace(u"\\N"_s, u"\n"_s).replace(u"\\n"_s, u"\n"_s).replace(u"\\h"_s, u" "_s);
|
||||
return normalizeLines(result);
|
||||
}
|
||||
|
||||
QList<SubtitleCue> mergeRepeated(QList<SubtitleCue> cues)
|
||||
{
|
||||
cues.removeIf([](const SubtitleCue &cue) {
|
||||
return cue.text.isEmpty() || cue.endMs <= cue.startMs;
|
||||
});
|
||||
std::stable_sort(cues.begin(), cues.end(), [](const SubtitleCue &a, const SubtitleCue &b) {
|
||||
return a.startMs < b.startMs;
|
||||
});
|
||||
|
||||
QList<SubtitleCue> merged;
|
||||
QHash<QString, qsizetype> lastIndexByText;
|
||||
for (const SubtitleCue &cue : std::as_const(cues)) {
|
||||
const auto it = lastIndexByText.constFind(cue.text);
|
||||
if (it != lastIndexByText.cend()) {
|
||||
SubtitleCue &previous = merged[*it];
|
||||
if (cue.startMs <= previous.endMs + MergeGapMs) {
|
||||
previous.endMs = std::max(previous.endMs, cue.endMs);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
lastIndexByText.insert(cue.text, merged.size());
|
||||
merged.append(cue);
|
||||
}
|
||||
merged.removeIf([](const SubtitleCue &cue) {
|
||||
return cue.endMs - cue.startMs < MinReadableMs;
|
||||
});
|
||||
return merged;
|
||||
}
|
||||
|
||||
QString toSrt(const QList<SubtitleCue> &cues)
|
||||
{
|
||||
QString out;
|
||||
for (qsizetype i = 0; i < cues.size(); ++i) {
|
||||
const SubtitleCue &cue = cues.at(i);
|
||||
out += QString::number(i + 1) + u'\n' + formatSrtTime(cue.startMs) + u" --> "_s + formatSrtTime(cue.endMs) + u'\n' + cue.text + u"\n\n"_s;
|
||||
}
|
||||
return out;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
#pragma once
|
||||
|
||||
#include <QList>
|
||||
#include <QString>
|
||||
|
||||
struct SubtitleCue {
|
||||
qint64 startMs = 0;
|
||||
qint64 endMs = 0;
|
||||
// Plain text without formatting; lines separated by '\n'.
|
||||
QString text;
|
||||
|
||||
bool operator==(const SubtitleCue &) const = default;
|
||||
};
|
||||
|
||||
namespace Subtitles
|
||||
{
|
||||
// SubRip and WebVTT.
|
||||
QList<SubtitleCue> parseSrt(const QString &content);
|
||||
// Dialogue events of an ASS/SSA script.
|
||||
QList<SubtitleCue> parseAss(const QString &content);
|
||||
// Picks the parser from the content. Sets \a error when the file can't be read.
|
||||
QList<SubtitleCue> parseFile(const QString &path, QString *error);
|
||||
|
||||
// Text of an ASS event without override tags. Empty for vector drawings.
|
||||
QString assEventText(QStringView text);
|
||||
|
||||
// Drops empty cues, sorts by start time, and merges cues that repeat the same text
|
||||
// back to back (typesetting often splits one line into many short frames). Cues that
|
||||
// stay too short to read after merging are dropped as animation frames.
|
||||
QList<SubtitleCue> mergeRepeated(QList<SubtitleCue> cues);
|
||||
|
||||
QString toSrt(const QList<SubtitleCue> &cues);
|
||||
}
|
||||
Reference in new issue
Block a user