import os
import argparse
import pathlib
import subprocess
import json
import time
import traceback
from concurrent.futures import ProcessPoolExecutor, as_completed
import aaf2
import opentimelineio as otio
from cgl.plugins.otio.tools import extract_shots
from cgl.plugins.otio.tools.aaf import aaf_embedded_media_tool, image_to_aaf
import cgl.plugins.alchemy.query as alc_query
FFMPEG_EXEC = alc_query.ffmpeg_path()
FFPROBE_EXEC = alc_query.ffprobe_path()
[docs]
def mp4_to_image_sequence(src, dest, audio_dst):
cmd = [FFMPEG_EXEC, '-y', '-nostdin']
cmd.extend(['-i', src, dest, audio_dst])
print(subprocess.list2cmdline(cmd))
subprocess.check_call(cmd)
[docs]
def image_list_to_mp4(src, dest):
cmd = [FFMPEG_EXEC, '-y', '-nostdin']
# cmd.extend(['-loglevel', 'warning'])
cmd.extend(['-f', 'concat', '-safe', '0', '-i', src])
# cmd.extend(['-r', '24', '-f', 'concat', '-safe', '0', '-i', src])
# cmd.extend(['-i', src])
# cmd.extend(['-acodec', 'copy'])
cmd.extend(['-vcodec', 'libx264', '-crf', '19', '-pix_fmt', 'yuv420p'])
cmd.extend([dest])
print(subprocess.list2cmdline(cmd))
subprocess.check_call(cmd)
[docs]
def get_master_mob_id(path):
try:
with aaf2.open(path, 'r') as f:
for mob in f.content.mastermobs():
return mob.mob_id
except:
return None
[docs]
def add_markers(aaf_path, markers):
with aaf2.open(aaf_path, 'rw') as f:
mob = next(f.content.toplevel())
video_slots = []
track_number = 1
for slot in mob.slots:
if slot.media_kind == 'Picture':
slot_id = slot["SlotID"].value
slot["PhysicalTrackNumber"].value = track_number
# print(track_number, slot_id)
video_slots.append([track_number, slot_id])
track_number +=1
event_slots = {}
for pos, track_index, marker in markers:
track_number, slot_id = video_slots[track_index]
m = f.create.DescriptiveMarker()
m['Position'].value = pos
m['Length'].value = 1
m['Comment'].value = marker.name
m['CommentMarkerUser'].value = "cgl"
m['DescribedSlots'].value = [slot_id]
comment = f.create.TaggedValue("Comment", marker.name)
m['UserComments'].append(comment)
if track_number not in event_slots:
event_slot = f.create.EventMobSlot()
event_slot.edit_rate = 24
event_slot.slot_id = 1000 + track_number
event_slot['PhysicalTrackNumber'].value = track_number
event_slot.segment = f.create.Sequence("DescriptiveMetadata")
mob.slots.append(event_slot)
event_slots[track_number] = event_slot
event_slots[track_number].segment.components.append(m)
[docs]
def get_audio_track(path):
mob_id = None
clip_name = None
length = None
with aaf2.open(path, 'r') as f:
mob = next(f.content.mastermobs())
mob_id = mob.mob_id
clip_name = mob.name
for slot in mob.slots:
length = slot.segment.length
break
assert mob_id
assert clip_name
assert length is not None and length > 0
clip = otio.schema.Clip(clip_name)
clip.source_range = otio.opentime.TimeRange(duration = otio.opentime.RationalTime(length, 24))
clip.metadata['AAF'] = {"SourceID": str(mob_id)}
media_ref = otio.schema.ExternalReference(
target_url = path,
available_range = clip.source_range
)
clip.media_reference = media_ref
track = otio.schema.Track(kind = otio.schema.TrackKind.Audio)
track.append(clip)
return track
[docs]
def get_video_track(mp4_path, aaf_path, mob_id):
clip_name = None
length = None
with aaf2.open(aaf_path, 'r') as f:
mob = f.content.mobs.get(mob_id)
clip_name = mob.name
for slot in mob.slots:
length = slot.segment.length
break
assert clip_name
assert length is not None and length > 0
clip = otio.schema.Clip(clip_name)
clip.source_range = otio.opentime.TimeRange(duration = otio.opentime.RationalTime(length, 24))
clip.metadata['AAF'] = {"SourceID": str(mob_id)}
media_ref = otio.schema.ExternalReference(
target_url = mp4_path,
available_range = clip.source_range
)
clip.media_reference = media_ref
track = otio.schema.Track(kind = otio.schema.TrackKind.Video)
track.append(clip)
return track
[docs]
def probe(path):
cmd = [FFPROBE_EXEC]
if path.lower().endswith(".txt"):
cmd.extend(["-f", "concat", "-safe", "0"])
cmd.extend(["-of", "json", "-show_format", "-show_streams", path])
# cmd.extend(["-of", "json", "-show_format", '-count_frames', "-show_streams", path])
print(subprocess.list2cmdline(cmd))
p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stdout, stderr = p.communicate()
if p.returncode != 0:
raise subprocess.CalledProcessError(
p.returncode, subprocess.list2cmdline(cmd), stderr
)
return json.loads(stdout)
[docs]
def still_dnxhr_encoder(image_files, output_dir):
txt_file = os.path.join(output_dir, "images.txt")
out_dnxhd = os.path.join(output_dir, "images.dnxhd")
files = image_files[:]
result = {}
bad_files = []
while files:
if os.path.exists(out_dnxhd):
os.remove(out_dnxhd)
with open(txt_file, 'w', encoding='utf8') as f:
f.write("ffconcat version 1.0\n")
for i, path in enumerate(files):
f.write(f"file '{path}'\n")
f.write(f"option framerate 24\n")
f.write(f"duration {1.0/24.0}\n")
f.write(f"""# {i}""" + "\n")
cmd = [FFMPEG_EXEC, '-y', '-nostdin']
cmd.extend(["-f", "concat", "-safe", "0"])
cmd.extend(['-i', txt_file])
width = 1920
height = 1080
cmd.extend(['-vf', f'scale=w={width}:h={height}:force_original_aspect_ratio=decrease,pad={width}:{height}:(ow-iw)/2:(oh-ih)/2'])
if width == 1920 and height == 1080:
# use dnxhd 36
cmd.extend(["-vcodec", "dnxhd", '-pix_fmt', 'yuv422p', '-vb', '36M'])
else:
cmd.extend(["-vcodec", "dnxhd", '-pix_fmt', 'yuv422p', "-profile:v", "dnxhr_lb"])
cmd.extend([out_dnxhd])
try:
subprocess.check_call(cmd)
except:
print(traceback.format_exc())
# if os.path.exists(out_dnxhd):
# os.remove(out_dnxhd)
current_files = {}
if os.path.exists(out_dnxhd):
with open(out_dnxhd, 'rb') as f:
for i,frame in enumerate(aaf2.video.iter_dnx_stream(f)):
# if i >= len(files):
# break
# print(i, len(files), len(frame))
src_image = files[i]
basename = os.path.basename(src_image)
convert_file = os.path.join(output_dir, basename + ".dnxhd")
with open(convert_file, 'wb') as f:
f.write(frame)
current_files[src_image] = convert_file
result[src_image] = convert_file
if len(current_files) != len(files):
files = files[len(current_files):]
bad_file = files.pop(0)
print(f"\nskipping {bad_file}\n")
bad_files.append(bad_file)
# random.shuffle(files)
else:
files = []
if bad_files:
bad_files.sort()
for file in bad_files:
print(f"bad file: {file}")
return result
[docs]
def mp4_chopper(mp4_path, edit_file_path, output_dir):
result = {}
start = time.time()
verify_shot_length = True
output_reference_aaf = True
is_aaf = extract_shots.is_aaf(edit_file_path)
timeline, shot_dict = extract_shots.simplify_timeline(edit_file_path)
probe_data = probe(mp4_path)
mp4_frame_count = None
for stream in probe_data['streams']:
if stream['codec_type'] != 'video':
continue
mp4_frame_count = int(stream['nb_frames'])
break
audio_channels = 0
for stream in probe_data['streams']:
if stream['codec_type'] != 'audio':
continue
channels = stream['channels']
audio_channels += channels
edit_basename = os.path.basename(edit_file_path)
edit_name, _ = os.path.splitext(edit_basename)
new_timeline = otio.schema.Timeline()
new_timeline.global_start_time = timeline.global_start_time
new_timeline.name = edit_basename
reference_aaf = os.path.join(output_dir, edit_basename + "_ref.aaf")
aaf_audio_files = []
for i in range(audio_channels):
audio_aaf = os.path.join(output_dir, f"audio_{edit_name}_{i:02d}.aaf")
aaf_audio_files.append(audio_aaf)
media_dict = {}
markers = []
for track_index, track in enumerate(timeline.video_tracks()):
new_track = otio.schema.Track()
new_timeline.tracks.append(new_track)
for item in track:
if isinstance(item, otio.schema.Transition):
# t = item.deepcopy()
# add_transition_metadata(t)
# new_track.append(t)
continue
psd_info = shot_dict.get(item.name, None)
if not psd_info:
new_track.append(item.deepcopy())
continue
# stack.append(item.deepcopy())
# print(stack)
new_item = item.deepcopy()
new_item = otio.algorithms.track_trimmed_to_range(new_item, item.source_range)
for clip in new_item.find_clips():
target_url = clip.media_reference.target_url
if target_url not in media_dict:
media_dict[target_url] = []
media_dict[target_url].append(clip)
shot_name = item.name
stack = otio.schema.Stack()
stack.append(new_item)
marker = otio.schema.Marker(shot_name)
stack.markers.append(marker)
new_track.append(stack)
assert isinstance(item, otio.schema.Track)
range_in_parent = item.range_in_parent()
cut_in = range_in_parent.start_time.to_frames()
cut_duration = range_in_parent.duration.to_frames()
markers.append([cut_in, track_index, marker])
assert cut_in >= 0
assert cut_in < mp4_frame_count
assert cut_in + cut_duration <= mp4_frame_count
full_dur = item.available_range().duration.to_frames()
head = item.source_range.start_time.to_frames()
tail = full_dur - cut_duration - head
print(f" {full_dur} {cut_in}-{cut_in+cut_duration} {head} {tail}")
shot_mp4 = os.path.join(output_dir, f"{shot_name}.mp4")
# shot_dnxhr = os.path.join(output_dir, f"{shot_name}.mov")
# shot_aaf = os.path.join(output_dir, f"{shot_name}.aaf")
shot_dnxhr = None
shot_aaf = None
result[shot_name] = { 'mp4_file' : shot_mp4}
if not os.path.exists(shot_mp4):
cmd = [FFMPEG_EXEC, '-y', '-nostdin']
cmd.extend(['-ss', str(cut_in/24.0), '-i', mp4_path])
cmd.extend(['-ss', str((cut_in+cut_duration-1)/24.0), '-i', mp4_path])
cmd.extend(['-f', 'lavfi', '-i', 'anullsrc=r=48000'])
filters = []
trimv = f'[0:v:0]trim=end={cut_duration/24.0}[trimv]'
trima = f'[0:a:0]atrim=end={cut_duration/24.0}[trima]'
filters.extend([trimv, trima])
if head:
# headv = f'[0:v:0]negate,trim=end_frame=1,loop=loop={head-1}:size=1[headv]'
headv = f'[0:v:0]trim=end_frame=1,loop=loop={head-1}:size=1[headv]'
heada = f'[2:a:0]atrim=end={head/24.0}[heada]'
filters.extend([headv, heada])
if tail:
# tailv = f'[1:v:0]negate,trim=end_frame=1,loop=loop={tail-1}:size=1[tailv]'
tailv = f'[1:v:0]trim=end_frame=1,loop=loop={tail-1}:size=1[tailv]'
taila = f'[2:a:0]atrim=end={tail/24.0}[taila]'
filters.extend([tailv, taila])
if head and tail:
concat = '[headv][heada][trimv][trima][tailv][taila]concat=n=3:v=1:a=1[outv][outa]'
elif head:
concat = '[headv][heada][trimv][trima]concat=n=2:v=1:a=1[outv][outa]'
elif tail:
concat = '[trimv][trima][tailv][taila]concat=n=2:v=1:a=1[outv][outa]'
else:
concat = '[trimv][trima]concat=n=1:v=1:a=1[outv][outa]'
filters.append(concat)
outv_h264 = '[outv]'
outa_h264 = '[outa]'
if shot_dnxhr:
vsplit = '[outv]split[outv_h264][outv_dnxhr]'
asplit = '[outa]asplit[outa_h264][outa_dnxhr]'
outv_h264 = '[outv_h264]'
outa_h264 = '[outa_h264]'
filters.extend([vsplit, asplit])
cmd.extend(["-filter_complex", ";".join(filters)])
cmd.extend(['-map', outv_h264, '-map', outa_h264])
cmd.extend(['-vcodec', 'libx264', '-crf', '19', '-pix_fmt', 'yuv420p'])
cmd.extend([shot_mp4])
if shot_dnxhr:
cmd.extend(['-map', "[outv_dnxhr]", '-map', "[outa_dnxhr]"])
cmd.extend(["-vcodec", "dnxhd", "-profile:v", "dnxhr_lb"])
cmd.extend([shot_dnxhr])
print(subprocess.list2cmdline(cmd))
subprocess.check_call(cmd)
if verify_shot_length:
# check the output ranges are correct
data = probe(shot_mp4)
for stream in data['streams']:
if stream['codec_type'] != 'video':
continue
# pprint(stream)
nb_frames = stream['nb_frames']
print(nb_frames, full_dur)
assert int(nb_frames) == full_dur
if shot_aaf:
if not os.path.exists(shot_aaf):
aaf_embedded_media_tool.create_aaf_file([shot_dnxhr],
shot_aaf,
item.name,
item.name,
frame_rate=24,
ignore_alpha = True,
use_embedded_timecode = True,
copy_dnxhd_streams = True
)
os.remove(shot_dnxhr)
mob_id = get_master_mob_id(shot_aaf)
media_ref = otio.schema.ExternalReference(
target_url = shot_aaf,
available_range = item.available_range()
)
clip = otio.schema.Clip(
item.name,
media_ref,
item.source_range
)
clip.metadata['AAF'] = {"SourceID": str(mob_id)}
new_track.append(clip)
result[shot_name]["aaf_file"] = shot_aaf
aaf_elements = []
aaf_sequence_file_path = None
if not is_aaf:
psd_files = []
for target_url, clips in media_dict.items():
if target_url.count("'"):
continue
basename = os.path.basename(target_url)
output_aaf = os.path.join(output_dir, f"{basename}.aaf")
# if os.path.exists(output_aaf):
# continue
name, ext = os.path.splitext(target_url)
if ext.lower() in ('.psd'):
psd_files.append(target_url)
dnxhr_mapping = still_dnxhr_encoder(psd_files, output_dir)
aaf_start = time.time()
if False:
for target_url, clips in media_dict.items():
basename = os.path.basename(target_url)
output_aaf = os.path.join(output_dir, f"{basename}.aaf")
dnxhr_file = dnxhr_mapping.get(target_url, None)
_, _, mob_id = process_psd_media(target_url, output_aaf, dnxhr_file)
for clip in clips:
clip.metadata['AAF'] = {"SourceID": str(mob_id)}
aaf_elements.append(output_aaf)
else:
executor = ProcessPoolExecutor(4)
try:
# process_psd_media(target_url, output_aaf, clips)
futures = []
for target_url, clips in media_dict.items():
basename = os.path.basename(target_url)
output_aaf = os.path.join(output_dir, f"{basename}.aaf")
dnxhr_file = dnxhr_mapping.get(target_url, None)
f = executor.submit(process_psd_media, target_url, output_aaf, dnxhr_file)
futures.append(f)
for f in as_completed(futures):
target_url, output_aaf, mob_id = f.result()
clips = media_dict[target_url]
for clip in clips:
clip.metadata['AAF'] = {"SourceID": str(mob_id)}
aaf_elements.append(output_aaf)
finally:
executor.shutdown(False, cancel_futures=True)
dur = time.time() - aaf_start
print(f"all aaf created in {dur} secs")
missing_audio_aaf_file = False
for path in aaf_audio_files:
if not os.path.exists(path):
missing_audio_aaf_file = True
break
if missing_audio_aaf_file:
wav_files = extract_audio_channels(mp4_path, output_dir, probe_data)
for aaf_path, wav_path in zip(aaf_audio_files, wav_files):
name = os.path.basename(wav_path)
aaf_embedded_media_tool.create_aaf_file([wav_path],
aaf_path,
name,
name,
frame_rate=24,
ignore_alpha = True,
use_embedded_timecode = True,
copy_dnxhd_streams = True
)
os.remove(wav_path)
for path in aaf_audio_files:
audio_track = get_audio_track(path)
new_timeline.tracks.append(audio_track)
aaf_elements.extend(aaf_audio_files)
if output_reference_aaf:
if not os.path.exists(reference_aaf):
name, _ = os.path.splitext(edit_basename)
aaf_embedded_media_tool.create_aaf_file([mp4_path],
reference_aaf,
name,
name,
frame_rate=24,
ignore_alpha = True,
use_embedded_timecode = True,
copy_dnxhd_streams = True
)
mob_id = get_master_mob_id(reference_aaf)
track = get_video_track(mp4_path, reference_aaf, mob_id)
new_timeline.tracks.insert(0, track)
aaf_elements.append(reference_aaf)
aaf_sequence_file_path = os.path.join(output_dir, f"{edit_basename}.aaf")
otio.adapters.write_to_file(new_timeline, aaf_sequence_file_path + ".otio")
otio.adapters.write_to_file(new_timeline, aaf_sequence_file_path, use_empty_mob_ids=True)
add_markers(aaf_sequence_file_path, markers)
sequence_aafs = [aaf_sequence_file_path]
dur = time.time() - start
print(f"mp4_chopper completed in {dur} secs")
return {"aaf_elements": aaf_elements, 'aaf_sequence': aaf_sequence_file_path, "shots" : result}
[docs]
def run_cli():
parser = argparse.ArgumentParser(
prog="chop timeline into multiple pieces based off a timeline",
)
parser.add_argument("-o", "--output-dir", type=pathlib.Path, default='.')
parser.add_argument("mp4_path", type=pathlib.Path)
parser.add_argument("xml_path", type=pathlib.Path)
args = parser.parse_args()
mp4_chopper(str(args.mp4_path), str(args.xml_path), str(args.output_dir))
if __name__ == "__main__":
run_cli()