From 9d7ecbf18aaa2b832912609445c147df28e4100d Mon Sep 17 00:00:00 2001 From: Vibhoothi Date: Sun, 19 Jul 2020 00:31:26 +0530 Subject: [PATCH] tools: draw-importance: Introduce --csv to export as csv Right now, when we do large processing of data, we are not storing the mean importance and frame type anywhere, so we require to save it to a file, by having an export option to csv, this makes things easier. To use: append --csv option along with other modes. This is only applicable to batch mode --- tools/draw-importances.py | 35 ++++++++++++++++++++++++----------- 1 file changed, 24 insertions(+), 11 deletions(-) diff --git a/tools/draw-importances.py b/tools/draw-importances.py index 0dfee25a2f..cb38c0e51a 100755 --- a/tools/draw-importances.py +++ b/tools/draw-importances.py @@ -20,9 +20,10 @@ @click.option('--verbose', is_flag=True, help="Will print verbose messages.") @click.option('--raw', is_flag=True, help="Print RAW Data of the bin.") @click.option('--figure', is_flag=True, help="Draw Importance to frame and save as PNG.") +@click.option('--csv', type = str, help="Export as frame_type, importance as CSV.") -def blockImportance(input, verbose, path, figure, raw): +def blockImportance(input, verbose, path, figure, raw, csv): """ CLI tool for extracting rav1e's Block Importance @@ -38,16 +39,28 @@ def blockImportance(input, verbose, path, figure, raw): Mode 2: Batch file processing for multiple hres and imps data\n Output files will be saved in out/YYYYMMDD_HH folder in -imps.png files. + Note: CSV Mode only available for Batch mode. """ + folder_path = "out/"+ datetime.datetime.now().strftime('%Y%m%d_%H') + if figure or csv: + if not os.path.exists('out'): + os.mkdir('out') + + if not os.path.exists(folder_path): + os.mkdir(folder_path) + if input == (None, None) and path == None : raise click.BadParameter("Please use either --input or --path method, for more information use --help.") if verbose: click.echo("Input given: "+ str(input)) + if path: click.echo("Path given: " + str(path)) if input != (None, None): + if csv: + raise click.BadParameter("Not Enough data to export as CSV, use Batch mode or remove --csv.") if verbose: click.secho("MODE 1", fg="yellow") with open(input[1], 'rb') as f: @@ -73,7 +86,7 @@ def blockImportance(input, verbose, path, figure, raw): col_list = [] frame_type_list = [] mean_list= [] - png_list = glob.glob(str(path)+"/*hres.png") + png_list = sorted(glob.glob(str(path)+"/*hres.png")) for png_iter in png_list: bin_list.append(png_iter.replace("hres.png","imps.bin")) total_files = len(bin_list) @@ -94,6 +107,11 @@ def blockImportance(input, verbose, path, figure, raw): click.secho("The full imps data after processing:", fg="red") click.echo(imps_list) + if csv: + np.savetxt(str(folder_path +'/'+ csv), np.column_stack( + (frame_type_list, mean_list)), delimiter=',', + header = "Frame Type, Mean Importance ", fmt = "%d, %f") + if verbose and path != None: click.secho("Frame Type List: "+ str(frame_type_list)) click.secho("\n png list: "+ str(png_list), fg="green") @@ -108,14 +126,6 @@ def blockImportance(input, verbose, path, figure, raw): frame_size_multiplier = 4 mv_original_block_size = 8 // 2 # The importances are in 8×8 blocks, but we use half-resolution images. mv_block_size = mv_original_block_size * frame_size_multiplier - folder_path = "out/"+ datetime.datetime.now().strftime('%Y%m%d_%H') - - if figure: - if not os.path.exists('out'): - os.mkdir('out') - - if not os.path.exists(folder_path): - os.mkdir(folder_path) if path == None and figure: frame = Image.open(input[0]) @@ -189,7 +199,10 @@ def blockImportance(input, verbose, path, figure, raw): if path == None: click.secho("File name is "+ splitext(input[1])[0]+ ".png", fg="blue") - print("Folder Path: "+ os.getcwd()+"/"+ folder_path) + print("Folder Path: "+ os.getcwd()+"/"+ folder_path) + + elif verbose and csv: + click.secho("CSV Saved sucessfully: "+ folder_path + '/'+ csv +".", fg="blue") click.secho("Processing done.", fg="green")