JVET_log.py

import argparse

import glob
import sys
import os


SEQUENCE_CLASSES = ['A1', 'A2', 'B', 'C', 'D', 'E', 'F']
SEQUENCE_PER_CLASSES = {
    'A1': ['Tango2', 'Drums100', 'Campfire', 'ToddlerFountain2'],
    'A2': ['CatRobot', 'TrafficFlow', 'DaylightRoad2', 'Rollercoaster2'],
    'B': ['Kimono', 'ParkScene', 'Cactus', 'BasketballDrive', 'BQTerrace'],
    'C': ['BasketballDrill', 'BQMall', 'PartyScene', 'RaceHorses'],
    'D': ['BasketballPass', 'BQSquare', 'BlowingBubbles', 'RaceHorses'],
    'E': ['FourPeople', 'Johnny', 'KristenAndSara'],
    'F': ['BasketballDrillText', 'ChinaSpeed', 'SlideEditing', 'SlideShow']
}
QPS = ['22', '27', '32', '37']


def extract_summary(file_name):
    f = open(file_name)
    while True:

        line = f.readline()
        if len(line) < 1:
            break
        
        if line.startswith('SUMMARY'):
            f.readline()
            summary = f.readline()
            summary = summary.split()

        elif 'Total Time' in line:
            total_time = line.split()[2]
            
        else:
            pass
    f.close()
    return summary[2:], total_time


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--log_dir')
    parser.add_argument('--output')
    args = parser.parse_args()

    log_dir = args.log_dir

    if log_dir == None:
        print('Usage : python JEM_log.py --log_dir <log directory path>')
        sys.exit()
        
    if log_dir not in os.listdir('.'):
        print('Invalid log directory')
        sys.exit()

    output_file = args.output
    if args.output == None:
        output_file = './log_summary.csv'
    


    log_list = glob.glob('./' + log_dir + '/*')

    

    f = open(output_file, 'w')
    f.write('file_name, Bitrate, Y-PSNR, U-PSNR, V-PSNR, Total Time\n')

    for log_file in log_list:
        if '\\' in log_file:
            file_name = log_file.split('\\')[-1]
        else:
            log_file.split('/')[-1]

        if '.' in file_name:
            file_name = file_name.split('.')[0]
        
        summary, total_time = extract_summary(log_file)
        result = [file_name] + summary[: -1] + [total_time]
        
        
        f.write(', '.join(result) + '\n')

    f.close()