194 lines
5.9 KiB
Python
Executable File
194 lines
5.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import lvsfunc as lvf
|
|
import typer
|
|
import vapoursynth as vs
|
|
import vsutil
|
|
|
|
core = vs.core
|
|
|
|
|
|
blackthr = 20
|
|
whitethr = 200
|
|
|
|
alphabet = (r" -!\"#%&'()*+,./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\]_"
|
|
"abcdefghijklmnopqrstuvwxyz{|}©«®°»ÀÂÆÇÈÉÊËÎÏÔÙÛÜàâäæçèéêëîïñôùûüÿŒœŸ‘’“”$£€fifl—…")
|
|
|
|
|
|
def invert(clip):
|
|
return core.std.Expr([clip], "x 255 - abs")
|
|
|
|
|
|
def diff_mask(hardsub, softsub):
|
|
diff_size = hardsub.width // 4, hardsub.height // 4
|
|
hy, hcb, hcr = vsutil.split(hardsub.resize.Point(*diff_size, format=vs.YUV444P8))
|
|
sy, scb, scr = vsutil.split(softsub.resize.Point(*diff_size, format=vs.YUV444P8))
|
|
diff_y = core.std.Expr([hy, sy], "x y - abs")
|
|
diff_cb = core.std.Expr([hcb, scb], "x y - abs")
|
|
diff_cr = core.std.Expr([hcr, scr], "x y - abs")
|
|
|
|
diff = core.std.Expr([diff_y, diff_cb, diff_cr], "x y z + +")
|
|
diff = diff.std.Binarize(40)
|
|
|
|
diff = diff.resize.Point(hardsub.width, hardsub.height)
|
|
return diff
|
|
|
|
|
|
# mask subs
|
|
def brightness_mask(hardsub):
|
|
white_sub = core.std.Binarize(vsutil.get_y(hardsub), whitethr)
|
|
black_sub = core.std.Binarize(vsutil.get_y(hardsub), blackthr)
|
|
black_sub = invert(black_sub)
|
|
white_expanded = vsutil.iterate(white_sub, core.std.Maximum, 6)
|
|
white_lightly_expanded = vsutil.iterate(white_sub, core.std.Maximum, 2)
|
|
|
|
matching_edges = core.std.Expr([white_expanded, black_sub], "x y and 255 *")
|
|
matching_edges = matching_edges.resize.Point(format=vs.GRAY8)
|
|
|
|
hyst_matching = core.misc.Hysteresis(matching_edges, white_lightly_expanded)
|
|
hyst_matching = core.std.MaskedMerge(core.std.BlankClip(hyst_matching),
|
|
hyst_matching, white_sub)
|
|
return matching_edges, hyst_matching
|
|
|
|
|
|
def filter_sub(hardsub, mask):
|
|
white_sub = core.std.Binarize(vsutil.get_y(hardsub), whitethr)
|
|
black = core.std.BlankClip(white_sub)
|
|
hardsub_expr = core.std.MaskedMerge(black, white_sub, mask, first_plane=True)
|
|
hardsub_expr = invert(hardsub_expr)
|
|
return hardsub_expr
|
|
|
|
|
|
def run_ocr(n, f, clip, ocr):
|
|
if f.props['HardsubDiffAverage'] > 0.0001:
|
|
# if f.props['_SceneChangePrev'] == 1 or f.props['_SceneChangeNext'] == 1 or n % 3 == 0:
|
|
return ocr
|
|
|
|
return clip
|
|
|
|
|
|
def set_sub_diff(n, f):
|
|
fout = f[0].copy()
|
|
for k, v in f[1].props.items():
|
|
if "HardsubDiff" in k:
|
|
fout.props[k] = v
|
|
|
|
return fout
|
|
|
|
|
|
def dump(n, f, ocr):
|
|
if (f.props.get('OCRString') or f.props['_SceneChangePrev'] == 1
|
|
or f.props['_SceneChangeNext'] == 1):
|
|
with open("ocr_dump.json", mode="a") as j:
|
|
props = {**f.props, "n": n}
|
|
if 'OCRString' in props:
|
|
props['OCRString'] = props['OCRString'].decode()
|
|
json.dump(props, j)
|
|
j.write("\n")
|
|
|
|
return ocr
|
|
|
|
|
|
def read_dump():
|
|
data = []
|
|
|
|
try:
|
|
with open("ocr_dump.json") as stream:
|
|
data = [json.loads(line) for line in stream]
|
|
except Exception as e:
|
|
print(e, file=sys.stderr)
|
|
|
|
to_list = lambda i: [i] if isinstance(i, int) else i
|
|
|
|
bad_ocr = sorted([i["n"] for i in data
|
|
if "OCRConfidence" in i and min(to_list(i["OCRConfidence"])) < 20])
|
|
|
|
|
|
def apply_ocr(clip):
|
|
ocr_clip = core.ocr.Recognize(clip, language="eng",
|
|
options=["enable_noise_removal", "false",
|
|
"tessedit_char_whitelist", alphabet,
|
|
"user_defined_dpi", "70"])
|
|
# ocr_clip = core.std.FrameEval(clip, functools.partial(run_ocr, clip=clip, ocr=ocr_clip), clip) # noqa
|
|
ocr_clip = ocr_clip.text.FrameProps()
|
|
return ocr_clip
|
|
|
|
|
|
def main(hardsub_file, raw_file, comparison=False):
|
|
src1 = lvf.src(hardsub_file)
|
|
src1 = src1.resize.Spline36(width=1920, height=1080)
|
|
# src1 = vsutil.depth(src1, 16)
|
|
src1 = src1.std.AssumeFPS(fpsnum=24000, fpsden=1001)
|
|
src2 = lvf.src(raw_file)
|
|
# src2 = vsutil.depth(src2, 16)
|
|
src2 = src2.std.AssumeFPS(fpsnum=24000, fpsden=1001)
|
|
|
|
sub_diff = diff_mask(vsutil.get_y(src1), vsutil.get_y(src2))
|
|
sub_diff = sub_diff.std.PlaneStats(prop="HardsubDiff")
|
|
|
|
sub_border_mask, sub_brightness = brightness_mask(src1)
|
|
|
|
hardsub_expr = filter_sub(src1, sub_diff)
|
|
hardsub_expr = hardsub_expr.std.PlaneStats()
|
|
hardsub_expr = core.std.ModifyFrame(hardsub_expr, [hardsub_expr, sub_diff], set_sub_diff)
|
|
hardsub_expr = core.misc.SCDetect(hardsub_expr, 0.0005)
|
|
|
|
ocr = apply_ocr(hardsub_expr)
|
|
brightness_ocr = apply_ocr(sub_brightness)
|
|
|
|
dumped_ocr = core.std.FrameEval(ocr, dump, ocr)
|
|
|
|
src1 = src1.std.PlaneStats()
|
|
src2 = src2.std.PlaneStats()
|
|
|
|
if comparison:
|
|
final = lvf.comparison.stack_vertical(
|
|
lvf.comparison.stack_horizontal(
|
|
sub_diff.resize.Point(format=src1.format),
|
|
brightness_ocr.resize.Point(format=src1.format),
|
|
src2.text.FrameProps()
|
|
),
|
|
lvf.comparison.stack_horizontal(
|
|
ocr.resize.Point(format=src1.format),
|
|
vsutil.get_y(src1).std.Binarize(whitethr).resize.Point(format=src1.format),
|
|
src1.text.FrameProps()
|
|
)
|
|
)
|
|
final = final.text.FrameNum(9)
|
|
final = final.text.FrameProps()
|
|
|
|
return final
|
|
else:
|
|
return dumped_ocr
|
|
|
|
|
|
app = typer.Typer()
|
|
|
|
|
|
@app.command()
|
|
def pipe(hardsub_file: Path, raw_file: Path):
|
|
final = main(hardsub_file, raw_file, comparison=True)
|
|
final.output(sys.stdout, y4m=True)
|
|
|
|
|
|
@app.command()
|
|
def mpv(hardsub_file: Path, raw_file: Path):
|
|
final = main(hardsub_file, raw_file, comparison=True)
|
|
mpv_proc = subprocess.Popen(["mpv", "-"], stdin=subprocess.PIPE)
|
|
final.output(mpv_proc.stdin, y4m=True)
|
|
mpv_proc.communicate()
|
|
|
|
|
|
@app.command(name="dump")
|
|
def dump_ocr():
|
|
pass
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app()
|