import os import sys # Maximum sprite size allowed by TB frontend, # see https://github.com/lanpa/tensorboardX/issues/516 TB_MAX_SPRITE_SIZE = 8192 def maybe_upload_file(local_path): '''Upload a file to remote cloud storage if the path starts with gs:// or s3:// ''' if local_path.startswith(('s3://', 'gs://')): prefix = local_path.split(':')[0] remote_bucket_path = local_path[len("s3://"):] # same length bp = remote_bucket_path.split("/") bucket = bp[0] path = remote_bucket_path[1 + len(bucket):] # s3://example/file becomes s3:/example/file in Linux local_path = prefix + ':/' + remote_bucket_path if prefix == 's3': import boto3 s3 = boto3.client('s3', endpoint_url=os.environ.get('S3_ENDPOINT')) s3.upload_file(local_path, bucket, path) elif prefix == 'gs': from google.cloud import storage client = storage.Client() Hbucket = storage.Bucket(client, bucket) blob = storage.Blob(path, Hbucket) blob.upload_from_filename(local_path) def make_tsv(metadata, save_path, metadata_header=None): if not metadata_header: metadata = [str(x) for x in metadata] else: assert len(metadata_header) == len(metadata[0]), \ 'len of header must be equal to the number of columns in metadata' metadata = ['\t'.join(str(e) for e in l) for l in [metadata_header] + metadata] named_path = os.path.join(save_path, 'metadata.tsv') if sys.version_info[0] == 3: with open(named_path, 'w', encoding='utf8') as f: for x in metadata: f.write(x + '\n') else: with open(named_path, 'wb') as f: for x in metadata: f.write((x + '\n').encode('utf-8')) maybe_upload_file(named_path) # https://github.com/tensorflow/tensorboard/issues/44 image label will be squared def make_sprite(label_img, save_path): import math import numpy as np from .x2num import make_np from .utils import make_grid from PIL import Image # this ensures the sprite image has correct dimension as described in # https://www.tensorflow.org/get_started/embedding_viz # There are some constraints for the sprite image: # 1. The sprite image should be square. # 2. Each image patch in the sprite image should be square. # 2. The content is row major order, so we can padding the image on the # bottom, but not on the right, otherwise, TB will treat some padded location # as images to be shown. # args: label_img: tensor in NCHW assert label_img.shape[2] == label_img.shape[3], 'Image should be square, see tensorflow/tensorboard#670' total_pixels = label_img.shape[0] * label_img.shape[2] * label_img.shape[3] pixels_one_side = total_pixels ** 0.5 number_of_images_per_row = int(math.ceil(pixels_one_side / label_img.shape[3])) arranged_img_CHW = make_grid(make_np(label_img), ncols=number_of_images_per_row) arranged_img_HWC = arranged_img_CHW.transpose(1, 2, 0) # chw -> hwc sprite_size = arranged_img_CHW.shape[2] assert sprite_size <= TB_MAX_SPRITE_SIZE, 'Sprite too large, see label_img shape limits' arranged_augment_square_HWC = np.ndarray((sprite_size, sprite_size, 3)) arranged_augment_square_HWC[:arranged_img_HWC.shape[0], :, :] = arranged_img_HWC im = Image.fromarray(np.uint8((arranged_augment_square_HWC * 255).clip(0, 255))) named_path = os.path.join(save_path, 'sprite.png') im.save(named_path) maybe_upload_file(named_path) def append_pbtxt(metadata, label_img, save_path, subdir, global_step, tag): from posixpath import join named_path = os.path.join(save_path, 'projector_config.pbtxt') with open(named_path, 'a') as f: # step = os.path.split(save_path)[-1] f.write('embeddings {\n') f.write('tensor_name: "{}:{}"\n'.format( tag, str(global_step).zfill(5))) f.write('tensor_path: "{}"\n'.format(join(subdir, 'tensors.tsv'))) if metadata is not None: f.write('metadata_path: "{}"\n'.format( join(subdir, 'metadata.tsv'))) if label_img is not None: f.write('sprite {\n') f.write('image_path: "{}"\n'.format(join(subdir, 'sprite.png'))) f.write('single_image_dim: {}\n'.format(label_img.shape[3])) f.write('single_image_dim: {}\n'.format(label_img.shape[2])) f.write('}\n') f.write('}\n') maybe_upload_file(named_path) def make_mat(matlist, save_path): named_path = os.path.join(save_path, 'tensors.tsv') with open(named_path, 'w') as f: for x in matlist: x = [str(i.item()) for i in x] f.write('\t'.join(x) + '\n') maybe_upload_file(named_path)