# import required libs
import glob
import os

import tensorflow as tf
import tensorflow_data_validation as tfdv
print('TF version: {}'.format(tf.version.VERSION))
print('TFDV version: {}'.format(tfdv.version.__version__))# Read artifact information from metadata store.
import beam_dag_runner

from tfx.orchestration import metadata
from tfx.types import standard_artifacts

metadata_connection_config = metadata.sqlite_metadata_connection_config(
              beam_dag_runner.METADATA_PATH)
with metadata.Metadata(metadata_connection_config) as store:
    stats_artifacts = store.get_artifacts_by_type(standard_artifacts.ExampleStatistics.TYPE_NAME)
    schema_artifacts = store.get_artifacts_by_type(standard_artifacts.Schema.TYPE_NAME)
    anomalies_artifacts = store.get_artifacts_by_type(standard_artifacts.ExampleAnomalies.TYPE_NAME)# configure output paths
# Exact paths to output artifacts can also be found on KFP Web UI if you are using kubeflow.
stats_path = stats_artifacts[-1].uri
train_stats_file = os.path.join(stats_path, 'train', 'stats_tfrecord')
eval_stats_file = os.path.join(stats_path, 'eval', 'stats_tfrecord')
print("Train stats file:{}, Eval stats file:{}".format(
    train_stats_file, eval_stats_file))

schema_file = os.path.join(schema_artifacts[-1].uri, 'schema.pbtxt')
print("Generated schame file:{}".format(schema_file))
anomalies_file = os.path.join(anomalies_artifacts[-1].uri, 'anomalies.pbtxt')
print("Generated anomalies file:{}".format(anomalies_file))# load generated statistics from StatisticsGen
train_stats = tfdv.load_statistics(train_stats_file)
eval_stats = tfdv.load_statistics(eval_stats_file)
tfdv.visualize_statistics(lhs_statistics=eval_stats, rhs_statistics=train_stats,
                          lhs_name='EVAL_DATASET', rhs_name='TRAIN_DATASET')# load generated schema from SchemaGen
schema = tfdv.load_schema_text(schema_file)
tfdv.display_schema(schema=schema)# load data vaildation result from ExampleValidator
anomalies = tfdv.load_anomalies_text(anomalies_file)
tfdv.display_anomalies(anomalies)