in python/vmaf/routine.py [0:0]
def read_dataset(dataset, **kwargs):
groundtruth_key = kwargs['groundtruth_key'] if 'groundtruth_key' in kwargs else None
skip_asset_with_none_groundtruth = kwargs['skip_asset_with_none_groundtruth'] \
if 'skip_asset_with_none_groundtruth' in kwargs else False
content_ids = kwargs['content_ids'] if 'content_ids' in kwargs else None
asset_ids = kwargs['asset_ids'] if 'asset_ids' in kwargs else None
workdir_root = kwargs['workdir_root'] if 'workdir_root' in kwargs else VmafConfig.workdir_path()
# asserts, can add more to the list...
assert hasattr(dataset, 'dataset_name')
assert hasattr(dataset, 'ref_videos')
assert hasattr(dataset, 'dis_videos')
assert hasattr(dataset, 'yuv_fmt') or all(['yuv_fmt' in ref_video for ref_video in dataset.ref_videos])
data_set_name = dataset.dataset_name
ref_videos = dataset.ref_videos
dis_videos = dataset.dis_videos
width = dataset.width if hasattr(dataset, 'width') else None
height = dataset.height if hasattr(dataset, 'height') else None
yuv_fmt = dataset.yuv_fmt if hasattr(dataset, 'yuv_fmt') else None
quality_width = dataset.quality_width if hasattr(dataset, 'quality_width') else None
quality_height = dataset.quality_height if hasattr(dataset, 'quality_height') else None
resampling_type = dataset.resampling_type if hasattr(dataset, 'resampling_type') else None
crop_cmd = dataset.crop_cmd if hasattr(dataset, 'crop_cmd') else None
pad_cmd = dataset.pad_cmd if hasattr(dataset, 'pad_cmd') else None
workfile_yuv_type = dataset.workfile_yuv_type if hasattr(dataset, 'workfile_yuv_type') else None
duration_sec = dataset.duration_sec if hasattr(dataset, 'duration_sec') else None
fps = dataset.fps if hasattr(dataset, 'fps') else None
start_frame = dataset.start_frame if hasattr(dataset, 'start_frame') else None
end_frame = dataset.end_frame if hasattr(dataset, 'end_frame') else None
ref_dict = {} # dictionary of content_id -> path for ref videos
for ref_video in ref_videos:
ref_dict[ref_video['content_id']] = ref_video
assets = []
for dis_video in dis_videos:
if content_ids is not None and dis_video['content_id'] not in content_ids:
continue
if asset_ids is not None and dis_video['asset_id'] not in asset_ids:
continue
if groundtruth_key is not None:
groundtruth = dis_video[groundtruth_key]
else:
if 'dmos' in dis_video:
groundtruth = dis_video['dmos']
elif 'mos' in dis_video:
groundtruth = dis_video['mos']
elif 'groundtruth' in dis_video:
groundtruth = dis_video['groundtruth']
else:
groundtruth = None
if 'os' in dis_video:
raw_groundtruth = dis_video['os']
else:
raw_groundtruth = None
if 'groundtruth_std' in dis_video:
groundtruth_std = dis_video['groundtruth_std']
else:
groundtruth_std = None
if 'rebuf_indices' in dis_video:
rebuf_indices = dis_video['rebuf_indices']
else:
rebuf_indices = None
ref_video = ref_dict[dis_video['content_id']]
ref_path = ref_video['path']
ref_yuv_fmt_ = yuv_fmt if yuv_fmt is not None else ref_dict[dis_video['content_id']]['yuv_fmt']
dis_yuv_fmt_ = dis_video['yuv_fmt'] if 'yuv_fmt' in dis_video else ref_yuv_fmt_
if width is not None:
width_ = width
elif 'width' in ref_video and 'width' not in dis_video:
width_ = ref_video['width']
elif 'width' in dis_video and 'width' not in ref_video:
width_ = dis_video['width']
elif 'width' in ref_video and 'width' in dis_video:
assert ref_video['width'] == dis_video['width']
width_ = ref_video['width']
else:
width_ = None
if height is not None:
height_ = height
elif 'height' in ref_video and 'height' not in dis_video:
height_ = ref_video['height']
elif 'height' in dis_video and 'height' not in ref_video:
height_ = dis_video['height']
elif 'height' in ref_video and 'height' in dis_video:
assert ref_video['height'] == dis_video['height']
height_ = ref_video['height']
else:
height_ = None
if quality_width is not None:
quality_width_ = quality_width
elif 'quality_width' in dis_video:
quality_width_ = dis_video['quality_width']
else:
quality_width_ = None
if quality_height is not None:
quality_height_ = quality_height
elif 'quality_height' in dis_video:
quality_height_ = dis_video['quality_height']
else:
quality_height_ = None
if resampling_type is not None:
resampling_type_ = resampling_type
elif 'resampling_type' in dis_video:
resampling_type_ = dis_video['resampling_type']
else:
resampling_type_ = None
if crop_cmd is not None:
ref_crop_cmd_ = crop_cmd
dis_crop_cmd_ = crop_cmd
else:
if 'crop_cmd' in ref_video:
ref_crop_cmd_ = ref_video['crop_cmd']
else:
ref_crop_cmd_ = None
if 'crop_cmd' in dis_video:
dis_crop_cmd_ = dis_video['crop_cmd']
else:
dis_crop_cmd_ = None
if pad_cmd is not None:
ref_pad_cmd_ = pad_cmd
dis_pad_cmd_ = pad_cmd
else:
if 'pad_cmd' in ref_video:
ref_pad_cmd_ = ref_video['pad_cmd']
else:
ref_pad_cmd_ = None
if 'pad_cmd' in dis_video:
dis_pad_cmd_ = dis_video['pad_cmd']
else:
dis_pad_cmd_ = None
if duration_sec is not None:
duration_sec_ = duration_sec
elif 'duration_sec' in dis_video:
duration_sec_ = dis_video['duration_sec']
else:
duration_sec_ = None
if fps is not None:
fps_ = fps
elif 'fps' in dis_video:
fps_ = dis_video['fps']
else:
fps_ = None
if start_frame is not None:
start_frame_ = start_frame
elif 'start_frame' in dis_video:
start_frame_ = dis_video['start_frame']
else:
start_frame_ = None
if end_frame is not None:
end_frame_ = end_frame
elif 'end_frame' in dis_video:
end_frame_ = dis_video['end_frame']
else:
end_frame_ = None
asset_dict = {'ref_yuv_type': ref_yuv_fmt_, 'dis_yuv_type': dis_yuv_fmt_}
if width_ is not None:
if asset_dict['ref_yuv_type'] != 'notyuv':
asset_dict['ref_width'] = width_
if asset_dict['dis_yuv_type'] != 'notyuv':
asset_dict['dis_width'] = width_
if height_ is not None:
if asset_dict['ref_yuv_type'] != 'notyuv':
asset_dict['ref_height'] = height_
if asset_dict['dis_yuv_type'] != 'notyuv':
asset_dict['dis_height'] = height_
if groundtruth is not None:
asset_dict['groundtruth'] = groundtruth
if raw_groundtruth is not None:
asset_dict['raw_groundtruth'] = raw_groundtruth
if groundtruth_std is not None:
asset_dict['groundtruth_std'] = groundtruth_std
if quality_width_ is not None:
asset_dict['quality_width'] = quality_width_
if quality_height_ is not None:
asset_dict['quality_height'] = quality_height_
if resampling_type_ is not None:
asset_dict['resampling_type'] = resampling_type_
if ref_crop_cmd_ is not None:
asset_dict['ref_crop_cmd'] = ref_crop_cmd_
if dis_crop_cmd_ is not None:
asset_dict['dis_crop_cmd'] = dis_crop_cmd_
if ref_pad_cmd_ is not None:
asset_dict['ref_pad_cmd'] = ref_pad_cmd_
if dis_pad_cmd_ is not None:
asset_dict['dis_pad_cmd'] = dis_pad_cmd_
if duration_sec_ is not None:
asset_dict['duration_sec'] = duration_sec_
if workfile_yuv_type is not None:
asset_dict['workfile_yuv_type'] = workfile_yuv_type
if rebuf_indices is not None:
asset_dict['rebuf_indices'] = rebuf_indices
if fps_ is not None:
asset_dict['fps'] = fps_
if start_frame_ is not None:
asset_dict['start_frame'] = start_frame_
if end_frame_ is not None:
asset_dict['end_frame'] = end_frame_
if 'ref_start_frame' in ref_video:
asset_dict['ref_start_frame'] = ref_video['ref_start_frame']
if 'dis_start_frame' in dis_video:
asset_dict['dis_start_frame'] = dis_video['dis_start_frame']
if 'ref_end_frame' in ref_video:
asset_dict['ref_end_frame'] = ref_video['ref_end_frame']
if 'dis_end_frame' in dis_video:
asset_dict['dis_end_frame'] = dis_video['dis_end_frame']
if 'enc_width' in dis_video:
asset_dict['dis_enc_width'] = dis_video['enc_width']
if 'enc_height' in dis_video:
asset_dict['dis_enc_height'] = dis_video['enc_height']
if 'enc_bitdepth' in dis_video:
asset_dict['dis_enc_bitdepth'] = dis_video['enc_bitdepth']
if groundtruth is None and skip_asset_with_none_groundtruth:
pass
else:
asset = Asset(dataset=data_set_name,
content_id=dis_video['content_id'],
asset_id=dis_video['asset_id'],
workdir_root=workdir_root,
ref_path=ref_path,
dis_path=dis_video['path'],
asset_dict=asset_dict,
)
assets.append(asset)
return assets