video2frame.py

import json
import re
import shutil
import subprocess
import warnings
from concurrent import futures
from pathlib import Path
from random import randint, random, shuffle

from tqdm import tqdm

from storage import STORAGE_TYPES
from util import parse_args, retry

ffmpeg_duration_template = re.compile(r"time=\s*(\d+):(\d+):(\d+)\.(\d+)")


def get_video_duration(video_file):
    cmd = [
        "ffmpeg",
        "-i", str(video_file),
        "-f", "null", "-"
    ]

    output = subprocess.check_output(cmd, stderr=subprocess.STDOUT)
    result_all = ffmpeg_duration_template.findall(output.decode())
    if result_all:
        result = result_all[-1]
        duration = float(result[0]) * 60 * 60 \
                   + float(result[1]) * 60 \
                   + float(result[2]) \
                   + float(result[3]) * (10 ** -len(result[3]))
    else:
        duration = -1
    return duration


def get_video_meta(video_file):
    try:
        cmd = [
            "ffprobe",
            "-v", "quiet",
            "-show_streams",
            "-print_format", "json",
            str(video_file)
        ]
        output = subprocess.check_output(cmd)
        output = json.loads(output)

        streamsbytype = {}
        for stream in output["streams"]:
            streamsbytype[stream["codec_type"].lower()] = stream

        return streamsbytype
    except:
        return {}


@retry()
def video_to_frames(args, video_file, video_meta, tmp_dir, error_when_empty=True):
    # Random clip the video
    clip_setting = []
    if args.duration > 0:
        try:
            video_duration = float(video_meta["video"]["duration"])
        except:
            video_duration = get_video_duration(video_file)
        if video_duration > 0:
            sta = max(0., random() * (video_duration - args.duration))
            dur = min(args.duration, video_duration - sta)
            clip_setting.extend([
                "-ss", "{}".format(sta),
                "-t", "{}".format(dur)
            ])
        else:
            warnings.warn("Ignore `duration` parameter for video {}.".format(video_file))

    cmd = [
        "ffmpeg",
        "-loglevel", "panic",
        "-vsync", "vfr",
        "-i", str(video_file),
        *args.vf_setting,
        *clip_setting,
        "-qscale:v", "2",
        str(tmp_dir / "%8d.jpg")
    ]
    subprocess.call(cmd)

    frames = [(int(f.name.split('.')[0]), f) for f in tmp_dir.iterdir()]
    frames.sort(key=lambda x: x[0])

    if error_when_empty and not frames:
        raise RuntimeError("Extract frame failed")

    return frames


@retry()
def sample_frames(args, frames, error_when_empty=True):
    if args.sample_mode:
        assert args.sample > 0, "Sample must >0, but get {}".format(args.sample)

        tot = len(frames)
        if args.sample_mode == 1:  # Uniformly sample n frames
            if args.sample == 1:
                index = [tot >> 1]
            else:
                step = (tot - 1.) / (args.sample - 1)
                index = [round(x * step) for x in range(args.sample)]
            frames = [frames[x] for x in index]
        elif args.sample_mode == 2:  # Randomly sample n continuous frames
            sta = randint(0, max(0, len(frames) - args.sample - 1))
            frames = frames[sta:sta + args.sample]
        elif args.sample_mode == 3:  # Randomly sample n frames
            shuffle(frames)
            frames = frames[:min(args.sample, tot)]
            frames.sort(key=lambda x: x[0])
        elif args.sample_mode == 4:  # Sample 1 frame every n frames
            frames = frames[::args.sample]
        else:
            raise AttributeError("Sample mode is not supported")

    if error_when_empty and not frames:
        raise RuntimeError("No frame selected")

    return frames


def process(args, video_key, video_info, frame_db):
    video_file = Path(video_info['path'])
    video_tmp_dir = Path(args.tmp_dir) / "{}".format(video_key)
    video_tmp_dir.mkdir(exist_ok=True)

    if not video_file.exists():
        raise RuntimeError("Video not exists")

    video_meta = get_video_meta(video_file)
    if not video_meta:
        raise RuntimeError("Can not get video info")

    for ith_clip in range(args.clips):
        clip_tmp_dir = video_tmp_dir / "{:03d}".format(ith_clip)
        clip_tmp_dir.mkdir(exist_ok=True, parents=True)

        # Get all frames
        frames = video_to_frames(args, video_file, video_meta, clip_tmp_dir)

        # Sample frames
        frames = sample_frames(args, frames)

        # Save to database
        frame_db.put(video_key, ith_clip, clip_tmp_dir, frames)

        if not args.keep:
            shutil.rmtree(clip_tmp_dir, ignore_errors=True)

    if not args.keep:
        shutil.rmtree(video_tmp_dir, ignore_errors=True)

    return "OK"


if "__main__" == __name__:
    args = parse_args()
    Path(args.tmp_dir).mkdir(exist_ok=True)

    frame_db = STORAGE_TYPES[args.db_type](args.db_name)

    annotation_all = json.load(Path(args.annotation_file).open())
    annotation = annotation_all["annotation"]
    total = len(annotation)
    fails = []

    if args.threads > 0:
        with futures.ThreadPoolExecutor(max_workers=args.threads) as executor:
            jobs = {
                executor.submit(process, args, video_key, video_info, frame_db): video_info['path']
                for video_key, video_info in annotation.items()
            }
            for future in tqdm(futures.as_completed(jobs), total=total):
                try:
                    video_status = future.result()
                except Exception as e:
                    tqdm.write("{} : {}".format(jobs[future], e))
                    fails.append(jobs[future])
                else:
                    tqdm.write("{} : {}".format(jobs[future], video_status))
    else:
        for video_key, video_info in tqdm(annotation.items()):
            try:
                video_status = process(args, video_key, video_info, frame_db)
            except Exception as e:
                tqdm.write("{} : {}".format(video_info['path'], e))
                fails.append(video_info['path'])
            else:
                tqdm.write("{} : {}".format(video_info['path'], video_status))

    frame_db.close()

    print("Processed {} videos".format(total))
    if not fails:
        print("All success! Congratulations!")
    else:
        print("{} Success, {} Error".format(total - len(fails), len(fails)))

        annotation = {k: v for k, v in annotation.items() if v['path'] not in fails}
        annotation_all["annotation"] = annotation
        if args.annotation_file.lower().endswith(".json"):
            save_path = args.annotation_file[:-5] + "-fix.json"
        else:
            save_path = args.annotation_file + "-fix.json"
        json.dump(annotation_all, Path(save_path).open("w"), indent=4)

    print("All Done!")