# Command-line interface for ColorVideoVDP. import sys import os.path import argparse import logging import glob import ffmpeg import numpy as np import torch import imageio.v2 as imageio import inspect import traceback from pycvvdp.vq_metric import vq_metric_dict import shlex import pycvvdp.utils as utils from pycvvdp.vq_exception import vq_exception from pycvvdp.display_model import vvdp_display_photometry, vvdp_display_geometry from pycvvdp.dump_channels import DumpChannels from pycvvdp.video_source_file import video_source_temp_resample_file, video_source_file # import all modules with `register_metric()` import pycvvdp.cvvdp_metric # noqa: F401 import pycvvdp.cvvdp_ml_metric # noqa: F401 import pycvvdp.dm_preview_metric # noqa: F401 import pycvvdp.psnr_metric # noqa: F401 import pycvvdp.vq_metric # noqa: F401 def expand_wildcards(filestrs): if not isinstance(filestrs, list): return [ filestrs ] files = [] for filestr in filestrs: if "*" in filestr: curlist = sorted( glob.glob(filestr) ) files = files + curlist else: files.append(filestr) return files # Save a numpy array as a video def np2vid(np_srgb, vidfile, fps, verbose=False): N, H, W, C = np_srgb.shape if C == 1: np_srgb = np.concatenate([np_srgb]*3, -1) process = ( ffmpeg .input('pipe:', format='rawvideo', pix_fmt='rgb24', s='{}x{}'.format(W, H), r=fps) #.output(vidfile, format='mp4', **{ "v:q": "10" } ) .output(vidfile, format='mp4', **{ "c:v": "mpeg4", "qscale:v": "3" } ) # mpeg4 codec is always bundled with ffmpeg so it should work .overwrite_output() .global_args( '-hide_banner') .global_args( '-loglevel', 'info' if verbose else 'quiet') .run_async(pipe_stdin=True) ) for fid in range(N): process.stdin.write( (np_srgb[fid,...] * 255.0) .astype(np.uint8) .tobytes() ) process.stdin.close() process.wait() # Save a numpy array as an image def np2img(np_srgb, imgfile): N, H, W, C = np_srgb.shape if C == 1: np_srgb = np.concatenate([np_srgb]*3, -1) if N>1: sys.exit("Expecting an image, found video") imageio.imwrite( imgfile, (np.clip(np_srgb,0.0,1.0)[0,...]*255.0).astype(np.uint8) ) # ----------------------------------- # Command-line Arguments # ----------------------------------- def parse_args(arg_list=None): available_metrics = [mm.replace('_', '-') for mm in vq_metric_dict.keys()] # Use dash in metric names parser = argparse.ArgumentParser(description="Evaluate ColorVideoVDP on a set of videos") parser.add_argument("-t", "--test", type=str, nargs='+', required = False, help="list of test images/videos") parser.add_argument("-r", "--ref", type=str, nargs='+', required = False, help="list of reference images/videos") parser.add_argument("--device", type=str, default='auto', help="select which PyTorch device to use. Pick from ['cpu', 'mps', 'cuda', 'cuda:0', 'cuda:1', ...]") parser.add_argument("--heatmap", type=str, default="none", help="type of difference map (none, raw, threshold, supra-threshold).") parser.add_argument("-g", "--distogram", type=float, default=-1, const=10, nargs='?', help="generate a distogram that visualizes the differences per-channel and per frame. The optional floating point parameter is the maximum JOD value to use in the visualization.") parser.add_argument("-x", "--features", action='store_true', default=False, help="generate JSON files with extracted features. Useful for retraining the metric.") parser.add_argument("-o", "--output-dir", type=str, default=None, help="in which directory heatmaps and feature files should be stored (the default is the current directory)") parser.add_argument("--result", type=str, default=None, help="write metric prediction results to a CSV file passed as an argument.") parser.add_argument("-c", "--config-paths", type=str, nargs='+', default=[], help="One or more paths to configuration files or directories. The main configurations files are `display_models.json`, `color_spaces.json` and `cvvdp_parameters.json`. The file name must start as the name of the original config file.") parser.add_argument("-d", "--display", type=str, default="standard_4k", help="display name, e.g. 'HTC Vive', or ? to print the list of models.") parser.add_argument("-n", "--nframes", type=int, default=-1, help="the number of video frames you want to compare") parser.add_argument("--count-frames", action='store_true', default=False, help="Use accurate method to count frames in a video. Slower but accurate. Use if you see frame read errors.") parser.add_argument("-f", "--full-screen-resize", choices=['bilinear', 'bicubic', 'nearest', 'area'], default=None, help="Both test and reference videos will be resized to match the full resolution of the display. Currently works only with videos.") parser.add_argument("-m", "--metric", choices=available_metrics, nargs='+', default=['cvvdp'], help='Select which metric(s) to run') parser.add_argument("--temp-padding", choices=['replicate', 'symmetric', 'valid'], default='replicate', help='How to pad the video in the time domain (for the temporal filters). "replicate" - repeat the first frame. "symmetric" - mirror the first frames. "valid" - skip initial frames until buffer is fully populated (no padding artifacts).') parser.add_argument("--pix-per-deg", type=float, default=None, help='Overwrite display geometry and use the provided pixels per degree value.') parser.add_argument("--fps", type=float, default=None, help='Frames per second. It will overwrite frame rate stores in the video file. Required when passing an array of image files.') parser.add_argument("--frames", type=str, default=None, help='Range of frames specified as first:step:last, first:last, or first: (Matlab notation). Currently works only with frames provided as images.') parser.add_argument("--gpu-mem", type=float, default=None, help='How much GPU memory can we use in GB. Use if CUDA reports out of mem errors, or you want to run multiple instances at the same time.') parser.add_argument("-q", "--quiet", action='store_true', default=False, help="Do not print any information but the final JOD value. Warning message will be still printed.") parser.add_argument("-v", "--verbose", action='store_true', default=False, help="Print out extra information.") parser.add_argument("--debug", action='store_true', default=False, help="Prints full stack trace when error is encountered.") parser.add_argument("--ffmpeg-cc", action='store_true', default=False, help="Use ffmpeg for upsampling and color conversion. Use custom pytorch code by default (faster and less memory).") parser.add_argument("--temp-resample", type=float, nargs="?", default=-1, const=0, help="Resample test and reference video to a common frame rate. Allows to compare videos of different frame rates. An optional argument - the maximum frame rate used when resampling.") parser.add_argument("-i", "--interactive", action='store_true', default=False, help="Run in an interactive mode, in which command line arguments are provided to the standard input, line by line. Saves on start-up time when running a large number of comparisons.") parser.add_argument("--dump-channels", nargs='+', choices=['temporal', 'lpyr', 'difference'], default=None, help="Output video/images with intermediate processing stages (for debugging and visualization).") if arg_list is not None: args = parser.parse_args(arg_list) else: args = parser.parse_args() return args def run_on_args(args): if args.quiet: log_level = logging.ERROR else: log_level = logging.DEBUG if args.verbose else logging.INFO logging.basicConfig(format='[%(levelname)s] %(message)s', level=log_level) args.metric = [mm.replace('-', '_') for mm in args.metric] # We need underscore for class names if args.verbose: import platform logging.debug( f'Platform: {platform.platform()}' ) if args.display == "?": vvdp_display_photometry.list_displays(args.config_paths) return if args.test is None or args.ref is None: logging.error( "Paths to both test and reference content needs to be specified.") return # Range of frames to process frame_range = None if not args.frames is None: ss = args.frames.split(':') if len(ss) == 3: sn = [0, 1, 10000] # default values else: sn = [0, 10000] for kk in range(len(ss)): if ss[kk].isnumeric(): sn[kk] = int(ss[kk]) if len(ss) == 3: frame_range = range(sn[0],(sn[2]+1),sn[1]) elif len(ss) <= 2: frame_range = range(sn[0],(sn[1]+1)) # Changed option to include MPS support for Macbooks # if args.gpu >= 0 and torch.cuda.is_available(): # device = torch.device('cuda:' + str(args.gpu)) # else: # device = torch.device('cpu') device = utils.get_best_device(args.device) # args.device = args.device.lower() # if args.device == 'auto': # Auto-detect CUDA or MPS # if torch.cuda.is_available(): # args.device = 'cuda' # elif torch.backends.mps.is_available(): # args.device = 'mps' # else: # logging.warning(f'No CUDA or MPS found and ColorVideoVDP will run on CPU. This may result in slow execution.') # args.device = 'cpu' # if args.device.startswith('cuda') and torch.cuda.is_available(): # device = torch.device(args.device) # elif args.device == 'mps': # torch_version = list(map(int, re.search(r'\d+\.\d+\.\d+', torch.__version__).group(0).split('.'))) # assert torch_version[0] > 2 or (torch_version[0] == 2 and torch_version[1] > 0), f'Please use torch>=2.1.0 with MPS. Current version is {torch.__version__}.' # assert sys.platform == 'darwin', 'Device "mps" is only valid on a Mac.' # device = torch.device(args.device) # else: # if args.device != 'cpu': # logging.warning(f'The requested device ({args.device}) is not found, reverting to CPU. This may result in slow execution.') # device = torch.device('cpu') logging.info("Running on device: " + str(device)) heatmap_types = ["raw", "threshold", "supra-threshold"] if args.heatmap == "none": args.heatmap = None if args.heatmap: if not args.heatmap in heatmap_types: logging.error( 'The recognized heatmap types are: "none", "raw", "threshold" and "supra-threshold"' ) sys.exit() do_heatmap = True else: do_heatmap = False # Check for valid resizing methods #if args.gpu_decode: # Doc: https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html #valid_methods = ['nearest', 'bilinear', 'bicubic', 'area', 'nearest-exact'] #else: # Doc: https://ffmpeg.org/ffmpeg-scaler.html #valid_methods = ['fast_bilinear', 'bilinear', 'bicubic', 'experimental', 'neighbor', 'area', 'bicublin', 'gauss', 'lanczos', 'spline'] #if args.full_screen_size not in valid_methods: # logging.error(f'The resizing method supplied is invalid. Please pick from {valid_methods}.') # sys.exit() args.test = expand_wildcards(args.test) args.ref = expand_wildcards(args.ref) N_test = len(args.test) N_ref = len(args.ref) if N_test==0: logging.error( "No test images/videos found." ) sys.exit() if N_ref==0: logging.error( "No reference images/videos found." ) sys.exit() if N_test != N_ref and N_test != 1 and N_ref != 1: logging.error( "Pass the same number of reference and test sources, or a single reference (to be used with all test sources), or a single test (to be used with all reference sources)." ) sys.exit() metrics = [] display_photometry = vvdp_display_photometry.load(args.display, config_paths=args.config_paths) if args.pix_per_deg is None: display_geometry = vvdp_display_geometry.load(args.display, config_paths=args.config_paths) else: display_geometry = vvdp_display_geometry( [1024, 1024], ppd=args.pix_per_deg ) out_dir = "." if args.output_dir is None else args.output_dir os.makedirs(out_dir, exist_ok=True) if args.dump_channels: dump_channels = DumpChannels( dump_temp_ch=("temporal" in args.dump_channels), dump_lpyr=("lpyr" in args.dump_channels), dump_diff=("difference" in args.dump_channels), output_dir=args.output_dir ) else: dump_channels = None for mm in args.metric: if not mm in vq_metric_dict: raise vq_exception( f"Unknown metric {mm}") metric_class = vq_metric_dict[mm] # The code below will figure out and pass only the parameters that a metric needs constructor_args = inspect.getfullargspec(metric_class.__init__)[0] all_bases = list(metric_class.__bases__) for bc in all_bases: constructor_args.extend( inspect.getfullargspec(bc.__init__)[0] ) all_bases.extend( list(bc.__bases__) ) met_args = {} if 'display_photometry' in constructor_args: met_args['display_photometry'] = display_photometry if 'display_geometry' in constructor_args: met_args['display_geometry'] = display_geometry if 'device' in constructor_args: met_args['device'] = device if 'heatmap' in constructor_args: met_args['heatmap'] = args.heatmap if 'temp_padding' in constructor_args: met_args['temp_padding'] = args.temp_padding if 'config_paths' in constructor_args: met_args['config_paths'] = args.config_paths if 'gpu_mem' in constructor_args: met_args['gpu_mem'] = args.gpu_mem if 'dump_channels' in constructor_args: met_args['dump_channels'] = dump_channels if 'quiet' in constructor_args: met_args['quiet'] = args.quiet fv = metric_class(**met_args) fv.train(False) metrics.append( fv ) info_str = metrics[-1].get_info_string() if not info_str is None: logging.info( 'When reporting metric results, please include the following information:' ) logging.info( info_str ) if not args.result is None: res_fh = open( args.result, "w" ) res_fh.write( 'test, reference' ) for mm in metrics: res_fh.write( ', ' + mm.short_name() ) res_fh.write( '\n' ) else: res_fh = None for kk in range( max(N_test, N_ref) ): # For each test and reference pair test_file = args.test[min(kk,N_test-1)] ref_file = args.ref[min(kk,N_ref-1)] if not res_fh is None: res_fh.write( f"{test_file}, {ref_file}" ) logging.info(f"Predicting the quality of '{test_file}' compared to '{ref_file}'") for mm in metrics: preload = False nframes = -2 if args.count_frames else args.nframes with torch.no_grad(): if args.temp_resample>=0: if args.temp_resample>0: video_source_temp_resample_file.max_fps = args.temp_resample vs = video_source_temp_resample_file( test_file, ref_file, display_photometry=display_photometry, config_paths=args.config_paths, full_screen_resize=args.full_screen_resize, resize_resolution=display_geometry.resolution, frames=nframes, ffmpeg_cc=args.ffmpeg_cc, verbose=args.verbose ) else: vs = video_source_file( test_file, ref_file, display_photometry=display_photometry, config_paths=args.config_paths, full_screen_resize=args.full_screen_resize, resize_resolution=display_geometry.resolution, frames=nframes, fps=args.fps, frame_range=frame_range, preload=preload, ffmpeg_cc=args.ffmpeg_cc, verbose=args.verbose ) base, ext = os.path.splitext(os.path.basename(test_file)) base_fname = os.path.join(out_dir, base) mm.set_base_fname(base_fname) Q_pred, stats = mm.predict_video_source(vs) Q_pred_scalar = Q_pred.item() if args.quiet: print( "{Q:0.4f}".format(Q=Q_pred_scalar) ) else: units_str = f" [{mm.quality_unit()}]" print( "{met_name}={Q:0.4f}{units}".format(met_name=mm.short_name(), Q=Q_pred_scalar, units=units_str) ) if not res_fh is None: res_fh.write( f", {Q_pred_scalar}" ) if args.features and not stats is None: if mm == 'pu-psnr': logging.warning( f'Skipping features as it is not supported by {mm}' ) break dest_name = os.path.join(out_dir, base + "_fmap.json") logging.info("Writing feature map '" + dest_name + "' ...") mm.write_features_to_json(stats, dest_name) if do_heatmap and not stats is None: if stats["heatmap"].shape[2]>1: # if it is a video dest_name = os.path.join(out_dir, base + "_heatmap.mp4") logging.info("Writing heat map '" + dest_name + "' ...") np2vid(torch.squeeze(stats["heatmap"].permute((2,3,4,1,0)), dim=4).cpu().numpy(), dest_name, vs.get_frames_per_second(), args.verbose) else: dest_name = os.path.join(out_dir, base + "_heatmap.png") logging.info("Writing heat map '" + dest_name + "' ...") np2img(torch.squeeze(stats["heatmap"].permute((2,3,4,1,0)), dim=4).cpu().numpy(), dest_name) if args.distogram != -1: dest_name = os.path.join(out_dir, base + "_distogram.png") logging.info("Writing distogram '" + dest_name + "' ...") jod_max = args.distogram mm.export_distogram( stats, dest_name, jod_max=jod_max ) del stats if not res_fh is None: res_fh.write( "\n" ) if not res_fh is None: res_fh.close() # del test_vid # torch.cuda.empty_cache() def main(): args = parse_args() try: if args.interactive: #print( "Running in an interactive mode" ) while True: line = sys.stdin.readline() if not line: break #print( shlex.split(line) ) args = parse_args(shlex.split(line)) run_on_args(args) else: run_on_args(args) except vq_exception as ex: logging.error( str(ex) ) if args.debug: traceback.print_exc() if __name__ == '__main__': main()