Skip to content

Commit

Permalink
update from main:repo.
Browse files Browse the repository at this point in the history
  • Loading branch information
endcloud committed Nov 11, 2022
1 parent ca2d2ed commit 93b84b1
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 20 deletions.
23 changes: 19 additions & 4 deletions autocut/main.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import argparse
import logging
import os

from autocut import utils

Expand All @@ -21,10 +22,12 @@ def main():
action=argparse.BooleanOptionalAction)
parser.add_argument('-s', help='Convert .srt to a compact format for easier editting',
action=argparse.BooleanOptionalAction)
parser.add_argument('-m', '--to-md', help='Convert .srt to .md for easier editting',
action=argparse.BooleanOptionalAction)
parser.add_argument('--lang', type=str, default='zh',
choices=['zh', 'en'],
help='The output language of transcription')
parser.add_argument('--prompt', type=str, default='大家好,',
parser.add_argument('--prompt', type=str, default='',
help='initial prompt feed into whisper')
parser.add_argument('--whisper-model', type=str, default='small',
choices=['tiny', 'base', 'small', 'medium', 'large'],
Expand All @@ -48,13 +51,25 @@ def main():
args = parser.parse_args()

if args.transcribe:
from autocut.transcribe import Transcribe
from .transcribe import Transcribe
Transcribe(args).run()
elif args.to_md:
from .utils import trans_srt_to_md
if len(args.inputs) == 2:
[input_1, input_2] = args.inputs
base, ext = os.path.splitext(input_1)
if ext != '.srt':
input_1, input_2 = input_2, input_1
trans_srt_to_md(args.encoding, args.force, input_1, input_2)
elif len(args.inputs) == 1:
trans_srt_to_md(args.encoding, args.force, args.inputs[0])
else:
logging.warn('Wrong number of files, please pass in a .srt file or an additional video file')
elif args.cut:
from autocut.cut import Cutter
from .cut import Cutter
Cutter(args).run()
elif args.daemon:
from autocut.daemon import Daemon
from .daemon import Daemon
Daemon(args).run()
elif args.s:
utils.compact_rst(args.inputs[0], args.encoding)
Expand Down
77 changes: 61 additions & 16 deletions autocut/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,23 @@
import os
import re

import srt
import opencc
import srt


def is_video(filename):
_, ext = os.path.splitext(filename)
return ext in ['.mp4', '.mov', '.mkv', '.flv']


def change_ext(filename, new_ext):
# Change the extension of filename to new_ext
base, _ = os.path.splitext(filename)
if not new_ext.startswith('.'):
new_ext = '.'+new_ext
new_ext = '.' + new_ext
return base + new_ext


def add_cut(filename):
# Add cut mark to the filename
base, ext = os.path.splitext(filename)
Expand All @@ -22,6 +28,7 @@ def add_cut(filename):
base += '_cut'
return base + ext


# a very simple markdown parser
class MD:
def __init__(self, filename, encoding) -> None:
Expand Down Expand Up @@ -66,14 +73,16 @@ def add_done_edditing(self, mark):

def add_video(self, video_fn):
ext = os.path.splitext(video_fn)[1][1:]
self.add(f'\n<video controls="true" allowfullscreen="true"> <source src="{video_fn}" type="video/{ext}"> </video>\n')
self.add(
f'\n<video controls="true" allowfullscreen="true"> <source src="{video_fn}" type="video/{ext}"> </video>\n')

def _parse_task_status(self, line):
# return (is_marked, rest) or (None, line) if not a task
m = re.match(r'- +\[([ x])\] +(.*)', line)
if not m:
return (None, line)
return (m.groups()[0].lower() == 'x', m.groups()[1])
return None, line
return m.groups()[0].lower() == 'x', m.groups()[1]


def check_exists(output, force):
if os.path.exists(output):
Expand All @@ -84,29 +93,32 @@ def check_exists(output, force):
return True
return False


def expand_segments(segments, expand_head, expand_tail, total_length):
# Pad head and tail for each time segment
results = []
for i in range(len(segments)):
t = segments[i]
start = max(t['start'] - expand_head,
segments[i-1]['end'] if i > 0 else 0)
segments[i - 1]['end'] if i > 0 else 0)
end = min(t['end'] + expand_tail,
segments[i+1]['start'] if i < len(segments)-1 else total_length)
results.append({'start':start, 'end':end})
segments[i + 1]['start'] if i < len(segments) - 1 else total_length)
results.append({'start': start, 'end': end})
return results


def remove_short_segments(segments, threshold):
# Remove segments whose length < threshold
return [s for s in segments if s['end'] - s['start'] > threshold]


def merge_adjacent_segments(segments, threshold):
# Merge two adjacent segments if their distance < threshold
results = []
i = 0
while i < len(segments):
s = segments[i]
for j in range(i+1, len(segments)):
for j in range(i + 1, len(segments)):
if segments[j]['start'] < s['end'] + threshold:
s['end'] = segments[j]['end']
i = j
Expand All @@ -116,6 +128,7 @@ def merge_adjacent_segments(segments, threshold):
results.append(s)
return results


def compact_rst(sub_fn, encoding):
cc = opencc.OpenCC('t2s')

Expand All @@ -131,17 +144,49 @@ def compact_rst(sub_fn, encoding):
subs = []
for l in lines:
items = l.split(' ')
if len(items) < 4: continue
if len(items) < 4:
continue
subs.append(srt.Subtitle(index=0,
start=srt.srt_timestamp_to_timedelta(items[0]),
end=srt.srt_timestamp_to_timedelta(items[2]),
content=' '.join(items[3:]).strip()))
with open(base[:-len(COMPACT)]+ext, 'wb') as f:
start=srt.srt_timestamp_to_timedelta(items[0]),
end=srt.srt_timestamp_to_timedelta(items[2]),
content=' '.join(items[3:]).strip()))
with open(base[:-len(COMPACT)] + ext, 'wb') as f:
f.write(srt.compose(subs).encode(encoding, 'replace'))
else:
# to a compact version
with open(sub_fn, encoding=encoding) as f:
subs = srt.parse(f.read())
with open(base+COMPACT+ext, 'wb') as f:
with open(base + COMPACT + ext, 'wb') as f:
for s in subs:
f.write(f'{srt.timedelta_to_srt_timestamp(s.start)} --> {srt.timedelta_to_srt_timestamp(s.end)} {cc.convert(s.content.strip())}\n'.encode(encoding, 'replace'))
f.write(
f'{srt.timedelta_to_srt_timestamp(s.start)} --> {srt.timedelta_to_srt_timestamp(s.end)} {cc.convert(s.content.strip())}\n'.encode(
encoding, 'replace'))


def trans_srt_to_md(encoding, force, srt_fn, video_fn=None):
base, ext = os.path.splitext(srt_fn)
if ext != '.srt':
logging.fatal('only .srt file is supported')
md_fn = base + ext.split(".")[0] + ".md"

check_exists(md_fn, force)

with open(srt_fn, encoding=encoding) as f:
subs = srt.parse(f.read())

md = MD(md_fn, encoding)
md.clear()
md.add_done_edditing(False)
if video_fn:
if not is_video(video_fn):
logging.fatal(f'{video_fn} may not be a video')
md.add_video(os.path.basename(video_fn))
md.add(f'\nTexts generated from [{os.path.basename(srt_fn)}]({os.path.basename(srt_fn)}).'
'Mark the sentences to keep for autocut.\n'
'The format is [subtitle_index,duration_in_second] subtitle context.\n\n')

for s in subs:
sec = s.start.seconds
pre = f'[{s.index},{sec // 60:02d}:{sec % 60:02d}]'
md.add_task(False, f'{pre:11} {s.content.strip()}')
md.write()

0 comments on commit 93b84b1

Please sign in to comment.