tomrunia · July 23, 2021 11:22
diff --git a/ffmpeg_utils.py b/ffmpeg_utils.py
 # MIT License
 #
 # Copyright (c) 2017 Tom Runia
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
 # in the Software without restriction, including without limitation the rights
 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 # copies of the Software, and to permit persons to whom the Software is
 # furnished to do so, subject to conditions.
 #
 # Author: Tom Runia
 # Date Created: 2017-04-07

 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function

 import os
 import shutil
 import math
 import subprocess as sp

 import cortex.utils

 ################################################################################
 ################################################################################
 ################################################################################

 '''
 The first part of this code borrows from: https://github.com/senko/python-data-converter/
 Mainly it contains helper classes to parse the output of 'ffprobe' to get data info.
 David Noelte et al. Copyright [C] 2011-2013. Python Video Converter contributors
 '''

 class MediaFormatInfo(object):
    """
    Describes the media container format. The attributes are:
      * format - format (short) name (eg. "ogg")
      * fullname - format full (descriptive) name
      * bitrate - total bitrate (bps)
      * duration - media duration in seconds
      * filesize - file size
    """

    def __init__(self):
        self.format = None
        self.fullname = None
        self.bitrate = None
        self.duration = None
        self.filesize = None
        self.num_frames = None

    def parse_ffprobe(self, key, val):
        """
        Parse raw ffprobe output (key=value).
        """
        if key == 'format_name':
            self.format = val
        elif key == 'format_long_name':
            self.fullname = val
        elif key == 'bit_rate':
            self.bitrate = MediaStreamInfo.parse_float(val, None)
        elif key == 'duration':
            self.duration = MediaStreamInfo.parse_float(val, None)
        elif key == 'nb_frames':
            self.num_frames = MediaStreamInfo.parse_int(val, None)
        elif key == 'size':
            self.size = MediaStreamInfo.parse_float(val, None)

    def __repr__(self):
        if self.duration is None:
            return 'MediaFormatInfo(format=%s)' % self.format
        return 'MediaFormatInfo(format=%s, duration=%.2f)' % (self.format,
                                                              self.duration)


 class MediaStreamInfo(object):
    """
    Describes one stream inside a media file. The general
    attributes are:
      * index - stream index inside the container (0-based)
      * type - stream type, either 'audio' or 'data'
      * codec - codec (short) name (e.g "vorbis", "theora")
      * codec_desc - codec full (descriptive) name
      * duration - stream duration in seconds
      * metadata - optional metadata associated with a data or audio stream
      * bitrate - stream bitrate in bytes/second
      * attached_pic - (0, 1 or None) is stream a poster image? (e.g. in mp3)
    Video-specific attributes are:
      * video_width - width of data in pixels
      * video_height - height of data in pixels
      * video_fps - average frames per second
    Audio-specific attributes are:
      * audio_channels - the number of channels in the stream
      * audio_samplerate - sample rate (Hz)
    """

    def __init__(self):
        self.index = None
        self.type = None
        self.codec = None
        self.codec_desc = None
        self.duration = None
        self.num_frames = None
        self.bitrate = None
        self.video_width = None
        self.video_height = None
        self.video_fps = None
        self.audio_channels = None
        self.audio_samplerate = None
        self.attached_pic = None
        self.sub_forced = None
        self.sub_default = None
        self.metadata = {}

    @staticmethod
    def parse_float(val, default=0.0):
        try:
            return float(val)
        except:
            return default

    @staticmethod
    def parse_int(val, default=0):
        try:
            return int(val)
        except:
            return default

    def parse_ffprobe(self, key, val):
        """
        Parse raw ffprobe output (key=value).
        """

        if key == 'index':
            self.index = self.parse_int(val)
        elif key == 'codec_type':
            self.type = val
        elif key == 'codec_name':
            self.codec = val
        elif key == 'codec_long_name':
            self.codec_desc = val
        elif key == 'duration':
            self.duration = self.parse_float(val)
        elif key == 'nb_frames':
            self.num_frames = self.parse_int(val)
        elif key == 'bit_rate':
            self.bitrate = self.parse_int(val, None)
        elif key == 'width':
            self.video_width = self.parse_int(val)
        elif key == 'height':
            self.video_height = self.parse_int(val)
        elif key == 'channels':
            self.audio_channels = self.parse_int(val)
        elif key == 'sample_rate':
            self.audio_samplerate = self.parse_float(val)
        elif key == 'DISPOSITION:attached_pic':
            self.attached_pic = self.parse_int(val)

        if key.startswith('TAG:'):
            key = key.split('TAG:')[1]
            value = val
            self.metadata[key] = value

        if self.type == 'audio':
            if key == 'avg_frame_rate':
                if '/' in val:
                    n, d = val.split('/')
                    n = self.parse_float(n)
                    d = self.parse_float(d)
                    if n > 0.0 and d > 0.0:
                        self.video_fps = float(n) / float(d)
                elif '.' in val:
                    self.video_fps = self.parse_float(val)

        if self.type == 'data':
            if key == 'r_frame_rate':
                if '/' in val:
                    n, d = val.split('/')
                    n = self.parse_float(n)
                    d = self.parse_float(d)
                    if n > 0.0 and d > 0.0:
                        self.video_fps = float(n) / float(d)
                elif '.' in val:
                    self.video_fps = self.parse_float(val)

        if self.type == 'subtitle':
            if key == 'disposition:forced':
                self.sub_forced = self.parse_int(val)
            if key == 'disposition:default':
                self.sub_default = self.parse_int(val)


    def __repr__(self):
        d = ''
        metadata_str = ['%s=%s' % (key, value) for key, value
                        in self.metadata.items()]
        metadata_str = ', '.join(metadata_str)

        if self.type == 'audio':
            d = 'type=%s, codec=%s, channels=%d, rate=%.0f' % (self.type,
                self.codec, self.audio_channels, self.audio_samplerate)
        elif self.type == 'data':
            d = 'type=%s, codec=%s, width=%d, height=%d, fps=%.1f' % (
                self.type, self.codec, self.video_width, self.video_height,
                self.video_fps)
        elif self.type == 'subtitle':
            d = 'type=%s, codec=%s' % (self.type, self.codec)
        if self.bitrate is not None:
            d += ', bitrate=%d' % self.bitrate

        if self.metadata:
            value = 'MediaStreamInfo(%s, %s)' % (d, metadata_str)
        else:
            value = 'MediaStreamInfo(%s)' % d

        return value


 class MediaInfo(object):
    """
    Information about media object, as parsed by ffprobe.
    The attributes are:
      * format - a MediaFormatInfo object
      * streams - a list of MediaStreamInfo objects
    """

    def __init__(self, posters_as_video=True):
        """
        :param posters_as_video: Take poster images (mainly for audio files) as
            A data stream, defaults to True
        """
        self._format = MediaFormatInfo()
        self.posters_as_video = posters_as_video
        self.streams = []

    def parse_ffprobe(self, raw):
        """
        Parse raw ffprobe output.
        """
        in_format = False
        current_stream = None

        for line in raw.split('\n'):
            line = line.strip()
            if line == '':
                continue
            elif line == '[STREAM]':
                current_stream = MediaStreamInfo()
            elif line == '[/STREAM]':
                if current_stream.type:
                    self.streams.append(current_stream)
                current_stream = None
            elif line == '[FORMAT]':
                in_format = True
            elif line == '[/FORMAT]':
                in_format = False
            elif '=' in line:
                k, v = line.split('=', 1)
                k = k.strip()
                v = v.strip()
                if current_stream:
                    current_stream.parse_ffprobe(k, v)
                elif in_format:
                    self._format.parse_ffprobe(k, v)

    @property
    def duration(self):
        return self._format.duration

    @property
    def num_frames(self):
        return self.streams[0].num_frames

    @property
    def format(self):
        return self._format.format

    @property
    def width(self):
        return self.streams[0].video_width

    @property
    def height(self):
        return self.streams[0].video_height

    @property
    def fps(self):
        return self.streams[0].video_fps

    @property
    def video_codec(self):
        for s in self.streams:
            if s.type == 'data':
                return s.codec
        return None

    def __repr__(self):
        return 'MediaInfo(format=%s, streams=%s)' % \
               (repr(self._format), repr(self.streams))


 ################################################################################
 ################################################################################
 ################################################################################

 def has_ffmpeg():
    if not shutil.which("ffmpeg"):  return False
    if not shutil.which("ffprobe"): return False
    return True

 def open_process(cmd):
    p = sp.Popen(cmd, shell=False, stdin=sp.PIPE, stdout=sp.PIPE, stderr=sp.PIPE, close_fds=True)
    return p

 def format_time(seconds):
    assert seconds <= 3600
    m = math.floor(seconds) // 60
    s = seconds-(m*60)
    return "00:{:02d}:{:.2f}".format(m,s)

 def media_info(video_file):
    '''
    Given an input video_file, calls 'ffprobe' to obtain data information. Then
     parses this output in a convenient MediaInfo object.
    
    :param video_file: str, data filename to get information on
    :return: MediaInfo object with ffprobe output
    '''

    if not os.path.exists(video_file):
        raise FileNotFoundError("Video file does not exist: {}".format(video_file))

    p = open_process(["ffprobe", "-show_format", "-show_streams", video_file])
    stdout_data, _ = p.communicate()
    stdout_data = stdout_data.decode("utf-8")

    # Parse ffprobe output
    info = MediaInfo()
    info.parse_ffprobe(stdout_data)

    if not info.format.format and len(info.streams) == 0:
        return None
    return info
	# MIT License
	#
	# Copyright (c) 2017 Tom Runia
	#
	# Permission is hereby granted, free of charge, to any person obtaining a copy
	# of this software and associated documentation files (the "Software"), to deal
	# in the Software without restriction, including without limitation the rights
	# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
	# copies of the Software, and to permit persons to whom the Software is
	# furnished to do so, subject to conditions.
	#
	# Author: Tom Runia
	# Date Created: 2017-04-07

	from __future__ import absolute_import
	from __future__ import division
	from __future__ import print_function

	import os
	import shutil
	import math
	import subprocess as sp

	import cortex.utils

	################################################################################
	################################################################################
	################################################################################

	'''
	The first part of this code borrows from: https://github.com/senko/python-data-converter/
	Mainly it contains helper classes to parse the output of 'ffprobe' to get data info.
	David Noelte et al. Copyright [C] 2011-2013. Python Video Converter contributors
	'''

	class MediaFormatInfo(object):
	"""
	Describes the media container format. The attributes are:
	* format - format (short) name (eg. "ogg")
	* fullname - format full (descriptive) name
	* bitrate - total bitrate (bps)
	* duration - media duration in seconds
	* filesize - file size
	"""

	def __init__(self):
	self.format = None
	self.fullname = None
	self.bitrate = None
	self.duration = None
	self.filesize = None
	self.num_frames = None

	def parse_ffprobe(self, key, val):
	"""
	Parse raw ffprobe output (key=value).
	"""
	if key == 'format_name':
	self.format = val
	elif key == 'format_long_name':
	self.fullname = val
	elif key == 'bit_rate':
	self.bitrate = MediaStreamInfo.parse_float(val, None)
	elif key == 'duration':
	self.duration = MediaStreamInfo.parse_float(val, None)
	elif key == 'nb_frames':
	self.num_frames = MediaStreamInfo.parse_int(val, None)
	elif key == 'size':
	self.size = MediaStreamInfo.parse_float(val, None)

	def __repr__(self):
	if self.duration is None:
	return 'MediaFormatInfo(format=%s)' % self.format
	return 'MediaFormatInfo(format=%s, duration=%.2f)' % (self.format,
	self.duration)


	class MediaStreamInfo(object):
	"""
	Describes one stream inside a media file. The general
	attributes are:
	* index - stream index inside the container (0-based)
	* type - stream type, either 'audio' or 'data'
	* codec - codec (short) name (e.g "vorbis", "theora")
	* codec_desc - codec full (descriptive) name
	* duration - stream duration in seconds
	* metadata - optional metadata associated with a data or audio stream
	* bitrate - stream bitrate in bytes/second
	* attached_pic - (0, 1 or None) is stream a poster image? (e.g. in mp3)
	Video-specific attributes are:
	* video_width - width of data in pixels
	* video_height - height of data in pixels
	* video_fps - average frames per second
	Audio-specific attributes are:
	* audio_channels - the number of channels in the stream
	* audio_samplerate - sample rate (Hz)
	"""

	def __init__(self):
	self.index = None
	self.type = None
	self.codec = None
	self.codec_desc = None
	self.duration = None
	self.num_frames = None
	self.bitrate = None
	self.video_width = None
	self.video_height = None
	self.video_fps = None
	self.audio_channels = None
	self.audio_samplerate = None
	self.attached_pic = None
	self.sub_forced = None
	self.sub_default = None
	self.metadata = {}

	@staticmethod
	def parse_float(val, default=0.0):
	try:
	return float(val)
	except:
	return default

	@staticmethod
	def parse_int(val, default=0):
	try:
	return int(val)
	except:
	return default

	def parse_ffprobe(self, key, val):
	"""
	Parse raw ffprobe output (key=value).
	"""

	if key == 'index':
	self.index = self.parse_int(val)
	elif key == 'codec_type':
	self.type = val
	elif key == 'codec_name':
	self.codec = val
	elif key == 'codec_long_name':
	self.codec_desc = val
	elif key == 'duration':
	self.duration = self.parse_float(val)
	elif key == 'nb_frames':
	self.num_frames = self.parse_int(val)
	elif key == 'bit_rate':
	self.bitrate = self.parse_int(val, None)
	elif key == 'width':
	self.video_width = self.parse_int(val)
	elif key == 'height':
	self.video_height = self.parse_int(val)
	elif key == 'channels':
	self.audio_channels = self.parse_int(val)
	elif key == 'sample_rate':
	self.audio_samplerate = self.parse_float(val)
	elif key == 'DISPOSITION:attached_pic':
	self.attached_pic = self.parse_int(val)

	if key.startswith('TAG:'):
	key = key.split('TAG:')[1]
	value = val
	self.metadata[key] = value

	if self.type == 'audio':
	if key == 'avg_frame_rate':
	if '/' in val:
	n, d = val.split('/')
	n = self.parse_float(n)
	d = self.parse_float(d)
	if n > 0.0 and d > 0.0:
	self.video_fps = float(n) / float(d)
	elif '.' in val:
	self.video_fps = self.parse_float(val)

	if self.type == 'data':
	if key == 'r_frame_rate':
	if '/' in val:
	n, d = val.split('/')
	n = self.parse_float(n)
	d = self.parse_float(d)
	if n > 0.0 and d > 0.0:
	self.video_fps = float(n) / float(d)
	elif '.' in val:
	self.video_fps = self.parse_float(val)

	if self.type == 'subtitle':
	if key == 'disposition:forced':
	self.sub_forced = self.parse_int(val)
	if key == 'disposition:default':
	self.sub_default = self.parse_int(val)


	def __repr__(self):
	d = ''
	metadata_str = ['%s=%s' % (key, value) for key, value
	in self.metadata.items()]
	metadata_str = ', '.join(metadata_str)

	if self.type == 'audio':
	d = 'type=%s, codec=%s, channels=%d, rate=%.0f' % (self.type,
	self.codec, self.audio_channels, self.audio_samplerate)
	elif self.type == 'data':
	d = 'type=%s, codec=%s, width=%d, height=%d, fps=%.1f' % (
	self.type, self.codec, self.video_width, self.video_height,
	self.video_fps)
	elif self.type == 'subtitle':
	d = 'type=%s, codec=%s' % (self.type, self.codec)
	if self.bitrate is not None:
	d += ', bitrate=%d' % self.bitrate

	if self.metadata:
	value = 'MediaStreamInfo(%s, %s)' % (d, metadata_str)
	else:
	value = 'MediaStreamInfo(%s)' % d

	return value


	class MediaInfo(object):
	"""
	Information about media object, as parsed by ffprobe.
	The attributes are:
	* format - a MediaFormatInfo object
	* streams - a list of MediaStreamInfo objects
	"""

	def __init__(self, posters_as_video=True):
	"""
	:param posters_as_video: Take poster images (mainly for audio files) as
	A data stream, defaults to True
	"""
	self._format = MediaFormatInfo()
	self.posters_as_video = posters_as_video
	self.streams = []

	def parse_ffprobe(self, raw):
	"""
	Parse raw ffprobe output.
	"""
	in_format = False
	current_stream = None

	for line in raw.split('\n'):
	line = line.strip()
	if line == '':
	continue
	elif line == '[STREAM]':
	current_stream = MediaStreamInfo()
	elif line == '[/STREAM]':
	if current_stream.type:
	self.streams.append(current_stream)
	current_stream = None
	elif line == '[FORMAT]':
	in_format = True
	elif line == '[/FORMAT]':
	in_format = False
	elif '=' in line:
	k, v = line.split('=', 1)
	k = k.strip()
	v = v.strip()
	if current_stream:
	current_stream.parse_ffprobe(k, v)
	elif in_format:
	self._format.parse_ffprobe(k, v)

	@property
	def duration(self):
	return self._format.duration

	@property
	def num_frames(self):
	return self.streams[0].num_frames

	@property
	def format(self):
	return self._format.format

	@property
	def width(self):
	return self.streams[0].video_width

	@property
	def height(self):
	return self.streams[0].video_height

	@property
	def fps(self):
	return self.streams[0].video_fps

	@property
	def video_codec(self):
	for s in self.streams:
	if s.type == 'data':
	return s.codec
	return None

	def __repr__(self):
	return 'MediaInfo(format=%s, streams=%s)' % \
	(repr(self._format), repr(self.streams))


	################################################################################
	################################################################################
	################################################################################

	def has_ffmpeg():
	if not shutil.which("ffmpeg"): return False
	if not shutil.which("ffprobe"): return False
	return True

	def open_process(cmd):
	p = sp.Popen(cmd, shell=False, stdin=sp.PIPE, stdout=sp.PIPE, stderr=sp.PIPE, close_fds=True)
	return p

	def format_time(seconds):
	assert seconds <= 3600
	m = math.floor(seconds) // 60
	s = seconds-(m*60)
	return "00:{:02d}:{:.2f}".format(m,s)

	def media_info(video_file):
	'''
	Given an input video_file, calls 'ffprobe' to obtain data information. Then
	parses this output in a convenient MediaInfo object.

	:param video_file: str, data filename to get information on
	:return: MediaInfo object with ffprobe output
	'''

	if not os.path.exists(video_file):
	raise FileNotFoundError("Video file does not exist: {}".format(video_file))

	p = open_process(["ffprobe", "-show_format", "-show_streams", video_file])
	stdout_data, _ = p.communicate()
	stdout_data = stdout_data.decode("utf-8")

	# Parse ffprobe output
	info = MediaInfo()
	info.parse_ffprobe(stdout_data)

	if not info.format.format and len(info.streams) == 0:
	return None
	return info