human-segmentation/python_core/media/VideoInfo.py

from distutils.log import fatal
import json

from ..common.Common import *
from ..common.Execute import *
from ..media.Timecodes import *

class VideoInfo:

    # file_path
    # file_size
    # video_codec: prores, h264, h265, vp9
    # width: Width of video, in pixels
    # height: Height of video, in pixels
    # color_space: ColorSpace of video
    # frame_rate: Frame rate of video, normally should be 24
    # start_timestamp: Starting timecode (PTS)
    # duration: Duration of video, in seconds
    # frame_count: Total frames in video
    # bit_depth: path of file

    def __init__(self, file_path):
        
        self.file_path = file_path
        self.file_size = GetFileSize(file_path)

        command = "ffprobe -v error -select_streams v:0 -show_entries stream=codec_name,height,width,r_frame_rate,color_space,start_time,duration,rotate,nb_frames -of csv=s=x:p=0 " + file_path
        stdout, _, errorCode = ExecuteCommand(command)

        stdout = stdout.split("\n")[0]
        stdout = stdout.strip()

        info = stdout.split("x")

        if len(info) < 8:
            log.error("File " + file_path + " has size " + str(self.Size))
            log.error("Error running command (return code: " + str(errorCode) +  "): " + command + " | Returned string is " + stdout + " | Array is " + str(info))

        index = 0

        self.video_codec = info[index]
        index = index + 1

        self.width = int(info[index])
        index = index + 1

        self.height = int(info[index])
        index = index + 1

        self.color_space = info[index]
        index = index + 1

        frame_rate_frac = info[index].split("/")
        index = index + 1

        self.frame_rate_float = float(frame_rate_frac[0])
        if len(frame_rate_frac) > 1:
            self.frame_rate_float = self.frame_rate_float / float(frame_rate_frac[1])
		
        self.frame_rate = round( self.frame_rate_float, 2 )

        if info[index] != "N/A":
            self.start_timestamp = float(info[index])
        index = index + 1

        if info[index] != "N/A":
            self.duration = float(info[index])
        index = index + 1

        if info[index] != "N/A":
            self.frame_count = int(info[index])
        else:
            self.frame_count = int(self.frame_rate_float * self.duration)

        if self.IsDolbyVision():
            self.bit_depth = 10
        else:
            self.bit_depth = 8

    def String(self):
            return json.dumps(self, default=lambda o: o.__dict__, 
                sort_keys=True, indent=4)

    def Print(self, stream):
        stream.write(self.String())

    def IsDolbyVision(self):
        return self.color_space == "bt2020nc"

    # Duration Getters
    def DurationAsMin(self):
        return self.duration / 60

    def DurationAsSec(self):
        return self.duration

    def DurationAsMillisec(self):
        return (self.duration * 1000)

    def DurationAsTimecode(self):
        return MillisecToTimeCode(self.DurationAsMillisec())