content: from __future__ import absolute_import
  from __future__ import print_function
  
  import json
  import os
  import pandas as pd
  from PIL import Image
  import shutil
  from six.moves.urllib.request import urlopen
  import subprocess
  import tempfile
  import unittest
  import uuid
  
  import tensorflow as tf
  from tensorflow.python.lib.io import file_io
  
  import google.datalab as dl
  import google.datalab.bigquery as bq
  import google.datalab.storage as storage
  
  CODE_PATH = os.path.abspath(os.path.join(
      os.path.dirname(__file__), '..', 'mltoolbox', 'code_free_ml'))
  
  # Some tests put files in GCS or use BigQuery. If HAS_CREDENTIALS is false,
  # those tests will not run.
  HAS_CREDENTIALS = True
  try:
    dl.Context.default().project_id
  except Exception:
    HAS_CREDENTIALS = False
  
  
  class TestTransformRawData(unittest.TestCase):
    """Tests for applying a saved model"""
  
    @classmethod
    def setUpClass(cls):
  
      # Set up dirs.
      cls.working_dir = tempfile.mkdtemp()
      cls.source_dir = os.path.join(cls.working_dir, 'source')
      cls.analysis_dir = os.path.join(cls.working_dir, 'analysis')
      cls.output_dir = os.path.join(cls.working_dir, 'output')
      file_io.create_dir(cls.source_dir)
  
      # Make test image files.
      img1_file = os.path.join(cls.source_dir, 'img1.jpg')
      image1 = Image.new('RGBA', size=(300, 300), color=(155, 0, 0))
      image1.save(img1_file)
      img2_file = os.path.join(cls.source_dir, 'img2.jpg')
      image2 = Image.new('RGBA', size=(50, 50), color=(125, 240, 0))
      image2.save(img2_file)
      img3_file = os.path.join(cls.source_dir, 'img3.jpg')
      image3 = Image.new('RGBA', size=(800, 600), color=(33, 55, 77))
      image3.save(img3_file)
  
      # Download inception checkpoint. Note that gs url doesn't work because
      # we may not have gcloud signed in when running the test.
      url = ('https://storage.googleapis.com/cloud-ml-data/img/' +
             'flower_photos/inception_v3_2016_08_28.ckpt')
      checkpoint_path = os.path.join(cls.working_dir, "checkpoint")
      response = urlopen(url)
      with open(checkpoint_path, 'wb') as f:
        f.write(response.read())
  
      # Make csv input file
      cls.csv_input_filepath = os.path.join(cls.source_dir, 'input.csv')
      file_io.write_string_to_file(
          cls.csv_input_filepath,
          '1,1,Monday,23.0,%s\n' % img1_file +
          '2,0,Friday,18.0,%s\n' % img2_file +
          '3,0,Sunday,12.0,%s\n' % img3_file)
  
      # Call analyze.py to create analysis results.
      schema = [{'name': 'key_col', 'type': 'INTEGER'},
                {'name': 'target_col', 'type': 'FLOAT'},
                {'name': 'cat_col', 'type': 'STRING'},
                {'name': 'num_col', 'type': 'FLOAT'},
                {'name': 'img_col', 'type': 'STRING'}]
      schema_file = os.path.join(cls.source_dir, 'schema.json')
      file_io.write_string_to_file(schema_file, json.dumps(schema))
      features = {'key_col': {'transform': 'key'},
                  'target_col': {'transform': 'target'},
                  'cat_col': {'transform': 'one_hot'},
                  'num_col': {'transform': 'identity'},
                  'img_col': {'transform': 'image_to_vec', 'checkpoint': checkpoint_path}}
      features_file = os.path.join(cls.source_dir, 'features.json')
      file_io.write_string_to_file(features_file, json.dumps(features))
      cmd = ['python ' + os.path.join(CODE_PATH, 'analyze.py'),
             '--output=' + cls.analysis_dir,
             '--csv=' + cls.csv_input_filepath,
             '--schema=' + schema_file,
             '--features=' + features_file]
      subprocess.check_call(' '.join(cmd), shell=True)
  
    @classmethod
    def tearDownClass(cls):
      shutil.rmtree(cls.working_dir)
  
    def test_local_csv_transform(self):
      """Test transfrom from local csv files."""
  
      cmd = ['python ' + os.path.join(CODE_PATH, 'transform.py'),
             '--csv=' + self.csv_input_filepath,
             '--analysis=' + self.analysis_dir,
             '--prefix=features',
             '--output=' + self.output_dir]
      print('cmd ', ' '.join(cmd))
      subprocess.check_call(' '.join(cmd), shell=True)
  
      # Read the tf record file. There should only be one file.
      record_filepath = os.path.join(self.output_dir,
                                     'features-00000-of-00001.tfrecord.gz')
      options = tf.python_io.TFRecordOptions(
          compression_type=tf.python_io.TFRecordCompressionType.GZIP)
      serialized_examples = list(tf.python_io.tf_record_iterator(record_filepath, options=options))
      self.assertEqual(len(serialized_examples), 3)
  
      # Find the example with key=1 in the file.
      first_example = None
      for ex in serialized_examples:
        example = tf.train.Example()
        example.ParseFromString(ex)
        if example.features.feature['key_col'].int64_list.value[0] == 1:
          first_example = example
      self.assertIsNotNone(first_example)
  
      transformed_number = first_example.features.feature['num_col'].float_list.value[0]
      self.assertAlmostEqual(transformed_number, 23.0)
  
      # transformed category = row number in the vocab file.
      transformed_category = first_example.features.feature['cat_col'].int64_list.value[0]
      vocab = pd.read_csv(
          os.path.join(self.analysis_dir, 'vocab_cat_col.csv'),
          header=None,
          names=['label', 'count'],
          dtype=str)
      origional_category = vocab.iloc[transformed_category]['label']
      self.assertEqual(origional_category, 'Monday')
  
      image_bytes = first_example.features.feature['img_col'].float_list.value
      self.assertEqual(len(image_bytes), 2048)
      self.assertTrue(any(x != 0.0 for x in image_bytes))
  
    @unittest.skipIf(not HAS_CREDENTIALS, 'GCS access missing')
    def test_local_bigquery_transform(self):
      """Test transfrom locally, but the data comes from bigquery."""
  
      # Make a BQ table, and insert 1 row.
      try:
        bucket_name = 'temp_pydatalab_test_%s' % uuid.uuid4().hex
        bucket_root = 'gs://%s' % bucket_name
        bucket = storage.Bucket(bucket_name)
        bucket.create()
  
        project_id = dl.Context.default().project_id
  
        dataset_name = 'test_transform_raw_data_%s' % uuid.uuid4().hex
        table_name = 'tmp_table'
  
        dataset = bq.Dataset((project_id, dataset_name)).create()
        table = bq.Table((project_id, dataset_name, table_name))
        table.create([{'name': 'key_col', 'type': 'INTEGER'},
                      {'name': 'target_col', 'type': 'FLOAT'},
                      {'name': 'cat_col', 'type': 'STRING'},
                      {'name': 'num_col', 'type': 'FLOAT'},
                      {'name': 'img_col', 'type': 'STRING'}])
  
        img1_file = os.path.join(self.source_dir, 'img1.jpg')
        dest_file = os.path.join(bucket_root, 'img1.jpg')
        file_io.copy(img1_file, dest_file)
  
        data = [
            {
             'key_col': 1,
             'target_col': 1.0,
             'cat_col': 'Monday',
             'num_col': 23.0,
             'img_col': dest_file,
            },
        ]
        table.insert(data=data)
  
        cmd = ['python ' + os.path.join(CODE_PATH, 'transform.py'),
               '--bigquery=%s.%s.%s' % (project_id, dataset_name, table_name),
               '--analysis=' + self.analysis_dir,
               '--prefix=features',
               '--project-id=' + project_id,
               '--output=' + self.output_dir]
        print('cmd ', ' '.join(cmd))
        subprocess.check_call(' '.join(cmd), shell=True)
  
        # Read the tf record file. There should only be one file.
        record_filepath = os.path.join(self.output_dir,
                                       'features-00000-of-00001.tfrecord.gz')
        options = tf.python_io.TFRecordOptions(
            compression_type=tf.python_io.TFRecordCompressionType.GZIP)
        serialized_examples = list(tf.python_io.tf_record_iterator(record_filepath, options=options))
        self.assertEqual(len(serialized_examples), 1)
  
        example = tf.train.Example()
        example.ParseFromString(serialized_examples[0])
  
        transformed_number = example.features.feature['num_col'].float_list.value[0]
        self.assertAlmostEqual(transformed_number, 23.0)
        transformed_category = example.features.feature['cat_col'].int64_list.value[0]
        self.assertEqual(transformed_category, 2)
        image_bytes = example.features.feature['img_col'].float_list.value
        self.assertEqual(len(image_bytes), 2048)
        self.assertTrue(any(x != 0.0 for x in image_bytes))
      finally:
        dataset.delete(delete_contents=True)
  
        for obj in bucket.objects():
          obj.delete()
        bucket.delete()
  
  
  if __name__ == '__main__':
      unittest.main()
  