mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-09 00:09:55 +00:00
Brandon Thomas
120 lines
3.8 KiB
Python
120 lines
3.8 KiB
Python
import argparse
|
|
import os
|
|
import subprocess
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
def get_video_resolution(input_video):
|
|
command = [
|
|
'ffprobe',
|
|
'-v', 'error',
|
|
'-select_streams', 'v:0',
|
|
'-show_entries', 'stream=width,height',
|
|
'-of', 'json',
|
|
input_video
|
|
]
|
|
|
|
print(f"Command: {command}", flush=True)
|
|
|
|
result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True)
|
|
video_info = json.loads(result.stdout)
|
|
|
|
width = video_info['streams'][0]['width']
|
|
height = video_info['streams'][0]['height']
|
|
|
|
return width, height
|
|
|
|
def calculate_new_resolution(width, height, max_resolution):
|
|
aspect_ratio = width / height
|
|
max_width, max_height = map(int, max_resolution.split('x'))
|
|
|
|
if width > max_width or height > max_height:
|
|
if (max_width / aspect_ratio) <= max_height:
|
|
new_width = max_width
|
|
new_height = round(new_width / aspect_ratio)
|
|
else:
|
|
new_height = max_height
|
|
new_width = round(new_height * aspect_ratio)
|
|
else:
|
|
new_width, new_height = width, height
|
|
|
|
new_width = 8 * round(new_width / 8)
|
|
new_height = 8 * round(new_height / 8)
|
|
|
|
return new_width, new_height
|
|
|
|
def format_timestamp(ms):
|
|
hours, ms = divmod(ms, 3600000)
|
|
minutes, ms = divmod(ms, 60000)
|
|
seconds, ms = divmod(ms, 1000)
|
|
return f"{int(hours):02d}:{int(minutes):02d}:{int(seconds):02d}.{int(ms):03d}"
|
|
|
|
def process_video(input_video, max_resolution, start_ms, end_ms, framerate, output_video = None):
|
|
width, height = get_video_resolution(input_video)
|
|
print(f"Original resolution: {width}x{height}")
|
|
|
|
# new_width, new_height = calculate_new_resolution(width, height, max_resolution)
|
|
# print(f"New resolution: {new_width}x{new_height}")
|
|
|
|
start_time = format_timestamp(start_ms)
|
|
duration_ms = end_ms - start_ms
|
|
duration = format_timestamp(duration_ms)
|
|
|
|
if not output_video:
|
|
# NB: This is salt's old implicit output video naming
|
|
input_vid_dir = Path(input_video).absolute().parent
|
|
output_video = input_vid_dir.joinpath(f"input.mp4")
|
|
|
|
command = [
|
|
'ffmpeg',
|
|
'-y',
|
|
'-i', input_video,
|
|
'-ss', start_time,
|
|
'-t', duration,
|
|
# '-vf', f'scale={new_width}:{new_height}',
|
|
'-r', str(framerate),
|
|
'-c:v', 'libx264', '-preset', 'fast', '-crf', '22',
|
|
'-c:a', 'aac', '-b:a', '128k',
|
|
output_video
|
|
]
|
|
|
|
print(f"Command: {command}", flush=True)
|
|
|
|
subprocess.run(command, check=True)
|
|
print(f"Video processed successfully. Output saved as {output_video}")
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 5:
|
|
# print("Usage: python script.py <input_video> <max_resolution> <start_ms> <end_ms> <framerate>")
|
|
print("Usage: python script.py <input_video> <start_ms> <end_ms> <framerate>")
|
|
sys.exit(1)
|
|
|
|
# Old Salt Args (replace this mess)
|
|
input_video = sys.argv[1]
|
|
output_video = None
|
|
# max_resolution = sys.argv[2]
|
|
max_resolution = None
|
|
start_ms = int(sys.argv[2])
|
|
end_ms = int(sys.argv[3])
|
|
framerate = float(sys.argv[4])
|
|
|
|
remaining_args = sys.argv[5:]
|
|
print('remaining args', remaining_args)
|
|
|
|
parser = argparse.ArgumentParser(description='Format Video')
|
|
# TODO: Output and input are optional for now until we migrate away from salt's arguments.
|
|
parser.add_argument('--input', type=str, help='Input Video Filename', required=False)
|
|
parser.add_argument('--output', type=str, help='Output Video Filename', required=False)
|
|
|
|
args = parser.parse_args(remaining_args)
|
|
|
|
if args.input:
|
|
input_video = args.input
|
|
|
|
if args.output:
|
|
output_video = args.output
|
|
|
|
process_video(input_video, max_resolution, start_ms, end_ms, framerate, output_video=output_video)
|