270 lines
11 KiB
JavaScript
270 lines
11 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];
|
|
const forced = stream.disposition && stream.disposition.forced;
|
|
const title = ((stream.tags && (stream.tags.title || stream.tags.TITLE)) || '').toLowerCase();
|
|
const isForced = Boolean(forced || title.includes('forced'));
|
|
if (isForced) parts.push('forced');
|
|
// Only ita.forced can be default. Do not add default to eng, non-forced, or other tracks.
|
|
if (lang === 'ita' && isForced) parts.push('default');
|
|
if (stream.disposition && stream.disposition.hearing_impaired) parts.push('sdh');
|
|
if (!parts.includes('sdh') && (title.includes('sdh') || title.includes('hearing impaired')
|
|
|| title.includes('hearing_impaired'))) 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) => normLang(l)).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 = [];
|
|
const taskByOutput = new Map();
|
|
|
|
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`);
|
|
|
|
const task = { stream, lang, codec, outSrt };
|
|
const isText = TEXT_CODECS.has(codec);
|
|
const isBitmap = BITMAP_CODECS.has(codec);
|
|
if (isText || isBitmap) {
|
|
const existing = taskByOutput.get(outSrt);
|
|
if (!existing || (isText && existing.type === 'bitmap')) {
|
|
taskByOutput.set(outSrt, { ...task, type: isText ? 'text' : 'bitmap' });
|
|
}
|
|
} else {
|
|
jobLog(`[extractSubtitlesAndOcr] Stream #${stream.index} codec "${codec}" not supported — skipping.\n`);
|
|
}
|
|
|
|
}
|
|
|
|
for (const task of taskByOutput.values()) {
|
|
if (task.type === 'text') textTasks.push(task);
|
|
else bitmapTasks.push(task);
|
|
}
|
|
|
|
// ── 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;
|