Choggbuster/ff_d2v.py

import sys
import json
import os
import subprocess as SP
import itertools as ITT
from tqdm import tqdm


colorspace={
    "gbr": 0,
    "bt709": 1,
    "unknown": 2,
    "fcc": 4,
    "bt470bg": 5,
    "smpte170m": 6,
    "smpte240m": 7,
    "ycgco": 8,
    "bt2020nc": 9,
    "bt2020c": 10,
    "smpte2085": 11,
    "chroma-derived-nc": 12,
    "chroma-derived-c": 13,
    "ictcp": 14,
}

pict_types={
    'I':0b01,
    'P':0b10,
    'B':0b11
}

def make_info(frames):
    has_interlaced = any(frame['interlaced_frame'] for frame in frames)
    new_gop='timecode' in frames[0].get('tags',{})
    info=0x000
    info|=1<<11 # always 1
    info|=0<<10 # 0=Closed GOP, 1=Open GOP
    info|=(not has_interlaced)<<9 # Progressive
    info|=new_gop<<8
    return info

def make_flags(frames):
    flags=[]
    for frame in frames:
        needs_prev=False
        progressive=not int(frame['interlaced_frame'])
        pt=pict_types[frame['pict_type']]
        reserved=0b00
        tff=int(frame['top_field_first'])
        rff=int(frame['repeat_pict'])
        flag=0b0
        flag|=(not needs_prev)<<7
        flag|=progressive<<6
        flag|=pt<<4
        flag|=reserved<<2
        flag|=tff<<1
        flag|=rff
        flags.append(f"{flag:02x}")
    return flags


def make_line(frames,stream):
    info=f"{make_info(frames):03x}"
    matrix=colorspace[stream['color_space']]
    file=0
    position=frames[0]['pkt_pos']
    skip=0
    vob=0
    cell=0
    flags=make_flags(frames)
    return " ".join(map(str,[info,matrix,file,position,skip,vob,cell,*flags]))


def get_frames(path):
    proc=SP.Popen([
        "ffprobe",
        "-probesize", str(0x7fffffff),
        "-analyzeduration", str(0x7fffffff),
        "-v","fatal",
        "-i",path,
        "-select_streams","v:0",
        "-show_frames",
        "-print_format","compact"
    ],stdout=SP.PIPE,stdin=SP.DEVNULL,bufsize=0)
    data=None
    for line in proc.stdout:
        line=str(line,"utf8").strip().split("|")
        line={line[0]: dict(v.split("=") for v in line[1:])}
        yield line
    ret=proc.wait()
    if ret!=0:
        exit(ret)
    return data

def get_streams(path):
    proc=SP.Popen([
        "ffprobe",
        "-probesize", str(0x7fffffff),
        "-analyzeduration", str(0x7fffffff),
        "-v","fatal",
        "-i",path,
        "-select_streams","v:0",
        "-show_streams",
        "-show_format",
        "-print_format","json"
    ],stdout=SP.PIPE,stdin=SP.DEVNULL,bufsize=0)
    data=json.load(proc.stdout)
    ret=proc.wait()
    if ret!=0:
        exit(ret)
    return data['streams'],data['format']

def make_header(file):
    return ["DGIndexProjectFile16","1",os.path.abspath(file)]

def make_settings(stream):
    pict_size="x".join(map(str,[stream["width"],stream["height"]]))
    frame_rate = list(map(int,stream['r_frame_rate'].split("/")))
    frame_rate=(frame_rate[0]*1000)//frame_rate[1]
    frame_rate=f"{frame_rate} ({stream['r_frame_rate']})"
    header=[
        ("Stream_Type",0), # Elementary Stream
        ("MPEG_Type",2), # MPEG-2
        ("iDCT_Algorithm",5), # 64-bit IEEE-1180 Reference
        ("YUVRGB_Scale",int(stream["color_range"]!="tv")),
        ("Luminance_Filter","0,0"),
        ("Clipping","0,0,0,0"),
        ("Aspect_Ratio",stream["display_aspect_ratio"]),
        ("Picture_Size",pict_size),
        ("Field_Operation",0), # Honor Pulldown Flags
        ("Frame_Rate", frame_rate),
        ("Location",f"0,0,0,0"),
    ]
    for k,v in header:
        yield f"{k}={v}"

def gen_d2v(path):
    yield from make_header(path)
    yield ""
    streams,fmt=get_streams(path)
    stream=[s for s in streams if s['codec_type']=='video'][0]
    stream['index']=str(stream['index'])
    yield from make_settings(stream)
    yield ""
    line_buffer=[]
    frames=get_frames(path)
    prog_bar=tqdm(frames,total=int(fmt['size']),unit_divisor=1024,unit_scale=True,unit="iB",dec="Writing d2v")
    for line in prog_bar:
        if 'frame' not in line:
            continue
        frame=line['frame']
        prog_bar.n=min(max(prog_bar.n,int(frame['pkt_pos'])),int(fmt['size']))
        prog_bar.update(0)
        if frame['stream_index']!=stream['index']:
            continue
        if frame['pict_type']=="I" and line_buffer:
            yield make_line(line_buffer,stream)
            line_buffer.clear()
        line_buffer.append(frame)
    prog_bar.close()
    yield None

def make_d2v(path):
    outfile=os.path.splitext(os.path.basename(path))[0]
    outfile=os.path.extsep.join([outfile,"d2v"])
    a,b=ITT.tee(gen_d2v(path))
    next(b)
    with open(outfile,"w") as fh:
        for line,next_line in zip(a,b):
            fh.write(line)
            if next_line is None: # last line, append end marker
                fh.write(" ff")
            fh.write("\n")