diff --git a/cbp_translate/components/subtitles.py b/cbp_translate/components/subtitles.py index d0e51d3..c289ccc 100644 --- a/cbp_translate/components/subtitles.py +++ b/cbp_translate/components/subtitles.py @@ -90,11 +90,21 @@ def add_subtitles( return np.array(img) +def get_border_h(height: int, border_size: float) -> int: + border_h = int(round(border_size * height)) + return border_h + + +def get_annotated_h(height: int, border_size: float) -> int: + border_h = get_border_h(height, border_size) + return height + border_h * 2 + + def add_borders(frame: Arr, size: float = 0.1) -> tuple[Arr, int]: """Add black borders to the top and bottom of a frame. That's where the subtitles will be displayed.""" height, width, _ = frame.shape - border_h = int(round(size * height)) + border_h = get_border_h(height, size) border = np.zeros((border_h, width, 3), np.uint8) frame = np.concatenate((border, frame, border), axis=0) diff --git a/cbp_translate/pipeline.py b/cbp_translate/pipeline.py index 874439b..9837f2f 100644 --- a/cbp_translate/pipeline.py +++ b/cbp_translate/pipeline.py @@ -27,6 +27,7 @@ add_borders, add_speaker_marker, add_subtitles, + get_annotated_h, ) from cbp_translate.components.translation import translate_segments from cbp_translate.modal_ import SHARED, cpu_image, deepl_secret, stub, volume @@ -171,6 +172,7 @@ def run(path_in: str, path_out: str, path_tmp: str, config: Config) -> Path: processed = annotate_frame.map(items, kwargs={"config": config}) # Now save the frames to an mp4 and add the original audio + height = get_annotated_h(height, config.border_size) save_frames(processed, fps, path_video, height=height, width=width) combine_streams(path_video, path_audio, path_out)