Files
Tdarr-plugin/LocalFlowPlugins/subtitles/extractSubtitlesAndOcr/1.0.0/index.js
T

255 lines
10 KiB
JavaScript

'use strict';
const path = require('path');
const fs = require('fs');
/* eslint-disable no-param-reassign */
/**
* extractSubtitlesAndOcr
* Extracts subtitle streams for ITA and ENG languages:
* - Text-based (subrip, srt, ass, ssa, webvtt):
* Fast extraction via mkvextract (all tracks in one pass, seconds instead of minutes),
* with single-pass FFmpeg (-vn -an) fallback.
* - Bitmap-based (hdmv_pgs_subtitle, dvd_subtitle): .sup extract
* then SubtitleEdit CLI OCR -> .srt (with graceful fallback on failure).
*
* Output files go to target directory (original library directory or next to source).
* Naming: <baseName>.<lang>[.forced][.default][.sdh].srt
*/
const TEXT_CODECS = new Set(['subrip', 'srt', 'ass', 'ssa', 'webvtt', 'mov_text']);
const BITMAP_CODECS = new Set(['hdmv_pgs_subtitle', 'dvd_subtitle', 'pgssub', 'dvbsub']);
const details = () => ({
name: 'Extract Subtitles & OCR Fallback',
description:
'Extracts ITA/ENG subtitle streams to SRT files in the library directory. '
+ 'Fast multi-track extraction via mkvextract / FFmpeg. '
+ 'Bitmap subtitles (PGS/VOBSUB) are OCR-d via SubtitleEdit CLI. '
+ 'OCR failures are logged as warnings without stopping the pipeline.',
style: { borderColor: '#00BCD4' },
tags: 'subtitles',
isStartPlugin: false,
pType: '',
requiresVersion: '2.11.01',
sidebarPosition: 3,
icon: 'faClosedCaptioning',
inputs: [
{
label: 'SubtitleEdit Executable Path',
name: 'subtitleEditPath',
type: 'string',
defaultValue: 'SubtitleEdit',
inputUI: { type: 'text' },
tooltip:
'Name or full path of the SubtitleEdit CLI executable. '
+ 'Must be in system PATH or provide absolute path. '
+ 'Used for OCR of PGS/VOBSUB streams on Windows workers.',
},
{
label: 'Languages to Extract (CSV)',
name: 'languages',
type: 'string',
defaultValue: 'ita,eng',
inputUI: { type: 'text' },
tooltip: 'Comma-separated list of ISO 639-2 language codes. Example: ita,eng',
},
],
outputs: [
{ number: 1, tooltip: 'Continue to next plugin (warnings in log if OCR failed)' },
],
});
// ─── Helpers ──────────────────────────────────────────────────────────────────
const normLang = (tag) => {
const l = (tag || '').toLowerCase().trim();
if (l === 'it') return 'ita';
if (l === 'en') return 'eng';
return l;
};
const buildSuffix = (lang, stream) => {
const parts = [lang];
if (stream.disposition && stream.disposition.forced) parts.push('forced');
if (stream.disposition && stream.disposition.default) parts.push('default');
if (stream.disposition && stream.disposition.hearing_impaired) parts.push('sdh');
const title = ((stream.tags && stream.tags.title) || '').toLowerCase();
if (!parts.includes('sdh') && title.includes('sdh')) parts.push('sdh');
return parts.join('.');
};
const getMkvExtract = (args) => {
let mkvExtract = 'mkvextract';
if (args.mkvpropeditPath) {
const dir = path.dirname(args.mkvpropeditPath);
const candidate = path.join(dir, process.platform === 'win32' ? 'mkvextract.exe' : 'mkvextract');
if (fs.existsSync(candidate)) return candidate;
}
return mkvExtract;
};
const runCLI = async (cliBin, spawnArgs, args) => {
const { CLI } = require('../../../../FlowHelpers/1.0.0/cliUtils');
const cli = new CLI({
cli: cliBin,
spawnArgs,
WorkDir: args.workDir,
jobLog: args.jobLog,
inputFileObj: args.inputFileObj,
updateWorker: args.updateWorker,
args,
});
return cli.runCli();
};
// ─── Plugin ───────────────────────────────────────────────────────────────────
const plugin = async (args) => {
const { jobLog, inputFileObj, variables, workDir } = args;
const inputFile = inputFileObj._id;
// Always save output SRTs in original library directory if known (e.g. from copyToWorkDir)
const targetDir = variables.originalLibraryDir || path.dirname(inputFile);
const origFile = variables.originalLibraryFile || inputFile;
const baseName = path.basename(origFile, path.extname(origFile));
jobLog(`[extractSubtitlesAndOcr] Input: ${inputFile}\n`);
jobLog(`[extractSubtitlesAndOcr] Output Directory: ${targetDir}\n`);
const subtitleEditPath = args.inputs.subtitleEditPath || 'SubtitleEdit';
const targetLangs = (args.inputs.languages || 'ita,eng')
.split(',').map((l) => l.trim().toLowerCase()).filter(Boolean);
const streams = (inputFileObj.ffProbeData && inputFileObj.ffProbeData.streams) || [];
const subtitleStreams = streams.filter((s) => s.codec_type === 'subtitle');
if (subtitleStreams.length === 0) {
jobLog('[extractSubtitlesAndOcr] No subtitle streams found.\n');
return { outputNumber: 1, outputFileObj: inputFileObj, variables };
}
jobLog(`[extractSubtitlesAndOcr] Found ${subtitleStreams.length} subtitle stream(s).\n`);
const textTasks = [];
const bitmapTasks = [];
for (const stream of subtitleStreams) {
const lang = normLang((stream.tags && stream.tags.language) || '');
if (!targetLangs.includes(lang)) {
jobLog(`[extractSubtitlesAndOcr] Stream #${stream.index} lang="${lang}" — not in target list, skipping.\n`);
continue;
}
const codec = (stream.codec_name || '').toLowerCase();
const suffix = buildSuffix(lang, stream);
const outSrt = path.join(targetDir, `${baseName}.${suffix}.srt`);
if (TEXT_CODECS.has(codec)) {
textTasks.push({ stream, lang, codec, outSrt });
} else if (BITMAP_CODECS.has(codec)) {
bitmapTasks.push({ stream, lang, codec, outSrt });
} else {
jobLog(`[extractSubtitlesAndOcr] Stream #${stream.index} codec "${codec}" not supported — skipping.\n`);
}
}
// ── 1. Batch extract text-based subtitles (mkvextract or single-pass FFmpeg)
if (textTasks.length > 0) {
jobLog(`[extractSubtitlesAndOcr] Extracting ${textTasks.length} text subtitle track(s)...\n`);
textTasks.forEach((t) => {
jobLog(` -> Stream #${t.stream.index} (${t.lang}, ${t.codec}) to ${path.basename(t.outSrt)}\n`);
});
let extracted = false;
const isMkv = inputFile.toLowerCase().endsWith('.mkv') || (inputFileObj.container === 'mkv');
const mkvExtract = getMkvExtract(args);
// Try mkvextract first (instantaneous on MKV)
if (isMkv) {
try {
const mkvArgs = ['tracks', inputFile];
textTasks.forEach((t) => mkvArgs.push(`${t.stream.index}:${t.outSrt}`));
const { cliExitCode } = await runCLI(mkvExtract, mkvArgs, args);
if (cliExitCode === 0) {
extracted = true;
jobLog('[extractSubtitlesAndOcr] mkvextract completed successfully.\n');
} else {
jobLog(`[extractSubtitlesAndOcr] mkvextract exited with code ${cliExitCode}, falling back to FFmpeg.\n`);
}
} catch (err) {
jobLog(`[extractSubtitlesAndOcr] mkvextract error (${err.message}), falling back to FFmpeg.\n`);
}
}
// Fallback: single-pass FFmpeg (-vn -an for fast demuxing without video/audio decoding)
if (!extracted) {
const ffmpegArgs = ['-y', '-i', inputFile, '-vn', '-an'];
textTasks.forEach((t, i) => {
ffmpegArgs.push('-map', `0:${t.stream.index}`, `-c:s:${i}`, 'srt', t.outSrt);
});
const { cliExitCode } = await runCLI(args.ffmpegPath, ffmpegArgs, args);
if (cliExitCode === 0) {
jobLog('[extractSubtitlesAndOcr] FFmpeg text extraction completed successfully.\n');
} else {
jobLog(`[extractSubtitlesAndOcr] WARNING: FFmpeg extraction failed with exit code ${cliExitCode}.\n`);
}
}
}
// ── 2. Process bitmap subtitles (PGS / VOBSUB) with SubtitleEdit OCR ─────
for (const t of bitmapTasks) {
const tmpSup = path.join(workDir, `${baseName}_s${t.stream.index}.sup`);
jobLog(`[extractSubtitlesAndOcr] Extracting bitmap stream #${t.stream.index} (${t.codec}) for OCR...\n`);
// Extract to .sup using mkvextract if mkv, else ffmpeg
let extractedSup = false;
const isMkv = inputFile.toLowerCase().endsWith('.mkv') || (inputFileObj.container === 'mkv');
if (isMkv) {
const mkvExtract = getMkvExtract(args);
const { cliExitCode } = await runCLI(mkvExtract, ['tracks', inputFile, `${t.stream.index}:${tmpSup}`], args);
if (cliExitCode === 0) extractedSup = true;
}
if (!extractedSup) {
const { cliExitCode } = await runCLI(
args.ffmpegPath,
['-y', '-i', inputFile, '-vn', '-an', '-map', `0:${t.stream.index}`, '-c:s', 'copy', tmpSup],
args,
);
if (cliExitCode !== 0) {
jobLog(`[extractSubtitlesAndOcr] WARNING: .sup extraction failed for stream #${t.stream.index}.\n`);
continue;
}
}
// Run OCR with SubtitleEdit CLI
try {
const seCli = await runCLI(
subtitleEditPath,
['/convert', tmpSup, 'subrip', '/track-number:0', `/outputfolder:${targetDir}`],
args,
);
if (seCli.cliExitCode !== 0) {
jobLog(`[extractSubtitlesAndOcr] WARNING: SubtitleEdit OCR failed (exit ${seCli.cliExitCode}) for stream #${t.stream.index}.\n`);
} else {
const seOut = path.join(targetDir, path.basename(tmpSup, '.sup') + '.srt');
if (fs.existsSync(seOut)) {
fs.renameSync(seOut, t.outSrt);
jobLog(`[extractSubtitlesAndOcr] OCR complete: ${path.basename(t.outSrt)}\n`);
}
}
} catch (err) {
jobLog(`[extractSubtitlesAndOcr] WARNING: SubtitleEdit not available: ${err.message}\n`);
} finally {
try { if (fs.existsSync(tmpSup)) fs.unlinkSync(tmpSup); } catch (_) { /* ignore */ }
}
}
jobLog('[extractSubtitlesAndOcr] Done.\n');
return { outputNumber: 1, outputFileObj: inputFileObj, variables };
};
module.exports.details = details;
module.exports.plugin = plugin;