def trim_video_frame()

in avhubert/preparation/lrs3_prepare.py [0:0]

• 41 lines of code
• 9 McCabe index (conditional complexity)

def trim_video_frame(csv_fn, raw_dir, output_dir, ffmpeg, rank, nshard):
    df = read_csv(csv_fn)
    raw2fid = OrderedDict()
    decimal, fps = 9, 25
    for fid, start, end in zip(df['id'], df['start'], df['end']):
        if '_' in fid:
            raw_fid = '_'.join(fid.split('_')[:-1])
        else:
            raw_fid = fid
        if raw_fid in raw2fid:
            raw2fid[raw_fid].append([fid, start, end])
        else:
            raw2fid[raw_fid] = [[fid, start, end]]
    i_raw = -1
    num_per_shard = math.ceil(len(raw2fid.keys())/nshard)
    start_id, end_id = num_per_shard*rank, num_per_shard*(rank+1)
    fid_info_shard = list(raw2fid.items())[start_id: end_id]
    print(f"Total videos in current shard: {len(fid_info_shard)}/{len(raw2fid.keys())}")
    for raw_fid, fid_info in tqdm(fid_info_shard):
        i_raw += 1
        raw_path = os.path.join(raw_dir, raw_fid+'.mp4')
        tmp_dir = tempfile.mkdtemp()
        cmd = ffmpeg + " -i " + raw_path + " " + tmp_dir + '/%0' + str(decimal) + 'd.png -loglevel quiet'
        subprocess.call(cmd, shell=True)
        num_frames = len(glob.glob(tmp_dir+'/*png'))
        for fid, start_sec, end_sec in fid_info:
            sub_dir = os.path.join(tmp_dir, fid)
            os.makedirs(sub_dir, exist_ok=True)
            start_sec, end_sec = float(start_sec), float(end_sec)
            if end_sec == -1:
                end_sec = 24*3600
            start_frame_id, end_frame_id = int(start_sec*fps), min(int(end_sec*fps), num_frames)
            imnames = [tmp_dir+'/'+str(x+1).zfill(decimal)+'.png' for x in range(start_frame_id, end_frame_id)]
            for ix, imname in enumerate(imnames):
                shutil.copyfile(imname, sub_dir+'/'+str(ix).zfill(decimal)+'.png')
            output_path = os.path.join(output_dir, fid+'.mp4')
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            cmd = [ffmpeg, "-i", sub_dir+'/%0'+str(decimal)+'d.png', "-y", "-crf", "20", output_path, "-loglevel", "quiet"]

            pipe = subprocess.call(cmd, stdout = subprocess.PIPE, stderr = subprocess.STDOUT) # subprocess.PIPE
        shutil.rmtree(tmp_dir)
    return