diff --git a/.vscode/.ropeproject/config.py b/.vscode/.ropeproject/config.py new file mode 100644 index 0000000..dee2d1a --- /dev/null +++ b/.vscode/.ropeproject/config.py @@ -0,0 +1,114 @@ +# The default ``config.py`` +# flake8: noqa + + +def set_prefs(prefs): + """This function is called before opening the project""" + + # Specify which files and folders to ignore in the project. + # Changes to ignored resources are not added to the history and + # VCSs. Also they are not returned in `Project.get_files()`. + # Note that ``?`` and ``*`` match all characters but slashes. + # '*.pyc': matches 'test.pyc' and 'pkg/test.pyc' + # 'mod*.pyc': matches 'test/mod1.pyc' but not 'mod/1.pyc' + # '.svn': matches 'pkg/.svn' and all of its children + # 'build/*.o': matches 'build/lib.o' but not 'build/sub/lib.o' + # 'build//*.o': matches 'build/lib.o' and 'build/sub/lib.o' + prefs['ignored_resources'] = ['*.pyc', '*~', '.ropeproject', + '.hg', '.svn', '_svn', '.git', '.tox'] + + # Specifies which files should be considered python files. It is + # useful when you have scripts inside your project. Only files + # ending with ``.py`` are considered to be python files by + # default. + # prefs['python_files'] = ['*.py'] + + # Custom source folders: By default rope searches the project + # for finding source folders (folders that should be searched + # for finding modules). You can add paths to that list. Note + # that rope guesses project source folders correctly most of the + # time; use this if you have any problems. + # The folders should be relative to project root and use '/' for + # separating folders regardless of the platform rope is running on. + # 'src/my_source_folder' for instance. + # prefs.add('source_folders', 'src') + + # You can extend python path for looking up modules + # prefs.add('python_path', '~/python/') + + # Should rope save object information or not. + prefs['save_objectdb'] = True + prefs['compress_objectdb'] = False + + # If `True`, rope analyzes each module when it is being saved. + prefs['automatic_soa'] = True + # The depth of calls to follow in static object analysis + prefs['soa_followed_calls'] = 0 + + # If `False` when running modules or unit tests "dynamic object + # analysis" is turned off. This makes them much faster. + prefs['perform_doa'] = True + + # Rope can check the validity of its object DB when running. + prefs['validate_objectdb'] = True + + # How many undos to hold? + prefs['max_history_items'] = 32 + + # Shows whether to save history across sessions. + prefs['save_history'] = True + prefs['compress_history'] = False + + # Set the number spaces used for indenting. According to + # :PEP:`8`, it is best to use 4 spaces. Since most of rope's + # unit-tests use 4 spaces it is more reliable, too. + prefs['indent_size'] = 4 + + # Builtin and c-extension modules that are allowed to be imported + # and inspected by rope. + prefs['extension_modules'] = [] + + # Add all standard c-extensions to extension_modules list. + prefs['import_dynload_stdmods'] = True + + # If `True` modules with syntax errors are considered to be empty. + # The default value is `False`; When `False` syntax errors raise + # `rope.base.exceptions.ModuleSyntaxError` exception. + prefs['ignore_syntax_errors'] = False + + # If `True`, rope ignores unresolvable imports. Otherwise, they + # appear in the importing namespace. + prefs['ignore_bad_imports'] = False + + # If `True`, rope will insert new module imports as + # `from import ` by default. + prefs['prefer_module_from_imports'] = False + + # If `True`, rope will transform a comma list of imports into + # multiple separate import statements when organizing + # imports. + prefs['split_imports'] = False + + # If `True`, rope will remove all top-level import statements and + # reinsert them at the top of the module when making changes. + prefs['pull_imports_to_top'] = True + + # If `True`, rope will sort imports alphabetically by module name instead + # of alphabetically by import statement, with from imports after normal + # imports. + prefs['sort_imports_alphabetically'] = False + + # Location of implementation of + # rope.base.oi.type_hinting.interfaces.ITypeHintingFactory In general + # case, you don't have to change this value, unless you're an rope expert. + # Change this value to inject you own implementations of interfaces + # listed in module rope.base.oi.type_hinting.providers.interfaces + # For example, you can add you own providers for Django Models, or disable + # the search type-hinting in a class hierarchy, etc. + prefs['type_hinting_factory'] = ( + 'rope.base.oi.type_hinting.factory.default_type_hinting_factory') + + +def project_opened(project): + """This function is called after opening the project""" + # Do whatever you like here! diff --git a/.vscode/.ropeproject/objectdb b/.vscode/.ropeproject/objectdb new file mode 100644 index 0000000..0a47446 Binary files /dev/null and b/.vscode/.ropeproject/objectdb differ diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 0000000..01adf0b --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,15 @@ +{ + // Use IntelliSense to learn about possible attributes. + // Hover to view descriptions of existing attributes. + // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 + "version": "0.2.0", + "configurations": [ + { + "name": "Python: Current File", + "type": "python", + "request": "launch", + "program": "${file}", + "console": "integratedTerminal" + } + ] +} diff --git a/README.md b/README.md index 1458044..1481cb9 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,18 @@ python test.py --clip_path ./sample_videos/clips/v32 \ Note that we use 216\*384 images for training, which has aspect ratio of 1:2. During inference, we scale the input to this size and then rescale the output back to the original size. +## Train + +We also provide training code for reference. The training can be started by running: + +```bash +python --data_root [root of video samples] \ + --data_root_imagenet [root of image samples] \ + --gpu_ids [gpu ids] \ +``` + +We do not provide the full video dataset due to the copyright issue. For image samples, we retrieve semantically similar images from ImageNet using [this repository](https://github.com/hmmlillian/Gray-Image-Retrieval). Still, one can refer to our code to understand the detailed procedure of augmenting the image dataset to mimic the video frames. + ## Comparison with State-of-the-Arts
@@ -73,6 +85,7 @@ Note that we use 216\*384 images for training, which has aspect ratio of 1:2. Du
## More results +
diff --git a/bash.sh b/bash.sh new file mode 100644 index 0000000..9aaca92 --- /dev/null +++ b/bash.sh @@ -0,0 +1,30 @@ +python test.py --clip_path ../test_videos/frames/0 --ref_path ../test_videos/refs/0 --output_path ../test_videos/output + +python test.py --clip_path ../test_videos/frames/1 --ref_path ../test_videos/refs/1 --output_path ../test_videos/output + +python test.py --clip_path ../test_videos/frames/2 --ref_path ../test_videos/refs/2 --output_path ../test_videos/output + +python test.py --clip_path ../test_videos/frames/3 --ref_path ../test_videos/refs/3 --output_path ../test_videos/output + +python test.py --clip_path ../test_videos/frames/4 --ref_path ../test_videos/refs/4 --output_path ../test_videos/output + +python test.py --clip_path ../test_videos/frames/5 --ref_path ../test_videos/refs/5 --output_path ../test_videos/output + +python test.py --clip_path ../test_videos/frames/6 --ref_path ../test_videos/refs/6 --output_path ../test_videos/output + +python test.py --clip_path ../test_videos/frames/7 --ref_path ../test_videos/refs/7 --output_path ../test_videos/output + +python test.py --clip_path ../test_videos/frames/8 --ref_path ../test_videos/refs/8 --output_path ../test_videos/output + +python test.py --clip_path ../test_videos/frames/9 --ref_path ../test_videos/refs/9 --output_path ../test_videos/output + +python test.py --clip_path ../test_videos/frames/10 --ref_path ../test_videos/refs/10 --output_path ../test_videos/output + +python test.py --clip_path ../test_videos/frames/11 --ref_path ../test_videos/refs/11 --output_path ../test_videos/output + +python test.py --clip_path ../test_videos/frames/12 --ref_path ../test_videos/refs/12 --output_path ../test_videos/output + +python test.py --clip_path ../test_videos/frames/13 --ref_path ../test_videos/refs/13 --output_path ../test_videos/output + +python test.py --clip_path ../test_videos/frames/14 --ref_path ../test_videos/refs/14 --output_path ../test_videos/output + diff --git a/lib/TrainTransforms.py b/lib/TrainTransforms.py new file mode 100644 index 0000000..1737927 --- /dev/null +++ b/lib/TrainTransforms.py @@ -0,0 +1,348 @@ +from __future__ import division + +import collections +import numbers +import random + +import torch +from PIL import Image +from skimage import color + +import lib.functional as F + +__all__ = [ + "Compose", + "Concatenate", + "ToTensor", + "Normalize", + "Resize", + "Scale", + "CenterCrop", + "Pad", + "RandomCrop", + "RandomHorizontalFlip", + "RandomVerticalFlip", + "RandomResizedCrop", + "RandomSizedCrop", + "FiveCrop", + "TenCrop", + "RGB2Lab", +] + + +def CustomFunc(inputs, func, *args, **kwargs): + im_l = func(inputs[0], *args, **kwargs) + im_ab = func(inputs[1], *args, **kwargs) + warp_ba = func(inputs[2], *args, **kwargs) + warp_aba = func(inputs[3], *args, **kwargs) + im_gbl_ab = func(inputs[4], *args, **kwargs) + bgr_mc_im = func(inputs[5], *args, **kwargs) + + layer_data = [im_l, im_ab, warp_ba, warp_aba, im_gbl_ab, bgr_mc_im] + + for l in range(5): + layer = inputs[6 + l] + err_ba = func(layer[0], *args, **kwargs) + err_ab = func(layer[1], *args, **kwargs) + + layer_data.append([err_ba, err_ab]) + + return layer_data + + +class Compose(object): + """Composes several transforms together. + + Args: + transforms (list of ``Transform`` objects): list of transforms to compose. + + Example: + >>> transforms.Compose([ + >>> transforms.CenterCrop(10), + >>> transforms.ToTensor(), + >>> ]) + """ + + def __init__(self, transforms): + self.transforms = transforms + + def __call__(self, inputs): + for t in self.transforms: + inputs = t(inputs) + return inputs + + +class Concatenate(object): + """ + Input: [im_l, im_ab, inputs] + inputs = [warp_ba_l, warp_ba_ab, warp_aba, err_pm, err_aba] + + Output:[im_l, err_pm, warp_ba, warp_aba, im_ab, err_aba] + """ + + def __call__(self, inputs): + im_l = inputs[0] + im_ab = inputs[1] + warp_ba = inputs[2] + warp_aba = inputs[3] + im_glb_ab = inputs[4] + bgr_mc_im = inputs[5] + bgr_mc_im = bgr_mc_im[[2, 1, 0], ...] + + err_ba = [] + err_ab = [] + + for l in range(5): + layer = inputs[6 + l] + err_ba.append(layer[0]) + err_ab.append(layer[1]) + + cerr_ba = torch.cat(err_ba, 0) + cerr_ab = torch.cat(err_ab, 0) + + return (im_l, cerr_ba, warp_ba, warp_aba, im_glb_ab, bgr_mc_im, im_ab, cerr_ab) + + +class ToTensor(object): + """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor. + + Converts a PIL Image or numpy.ndarray (H x W x C) in the range + [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0]. + """ + + def __call__(self, inputs): + """ + Args: + pic (PIL Image or numpy.ndarray): Image to be converted to tensor. + + Returns: + Tensor: Converted image. + """ + return CustomFunc(inputs, F.to_mytensor) + + +class Normalize(object): + """Normalize an tensor image with mean and standard deviation. + Given mean: ``(M1,...,Mn)`` and std: ``(S1,..,Sn)`` for ``n`` channels, this transform + will normalize each channel of the input ``torch.*Tensor`` i.e. + ``input[channel] = (input[channel] - mean[channel]) / std[channel]`` + + Args: + mean (sequence): Sequence of means for each channel. + std (sequence): Sequence of standard deviations for each channel. + """ + + def __call__(self, inputs): + """ + Args: + tensor (Tensor): Tensor image of size (C, H, W) to be normalized. + + Returns: + Tensor: Normalized Tensor image. + """ + + im_l = F.normalize(inputs[0], 50, 1) # [0, 100] + im_ab = F.normalize(inputs[1], (0, 0), (1, 1)) # [-100, 100] + + inputs[2][0:1, :, :] = F.normalize(inputs[2][0:1, :, :], 50, 1) + inputs[2][1:3, :, :] = F.normalize(inputs[2][1:3, :, :], (0, 0), (1, 1)) + warp_ba = inputs[2] + + inputs[3][0:1, :, :] = F.normalize(inputs[3][0:1, :, :], 50, 1) + inputs[3][1:3, :, :] = F.normalize(inputs[3][1:3, :, :], (0, 0), (1, 1)) + warp_aba = inputs[3] + + im_gbl_ab = F.normalize(inputs[4], (0, 0), (1, 1)) # [-100, 100] + + bgr_mc_im = F.normalize(inputs[5], (123.68, 116.78, 103.938), (1, 1, 1)) + + layer_data = [im_l, im_ab, warp_ba, warp_aba, im_gbl_ab, bgr_mc_im] + + for l in range(5): + layer = inputs[6 + l] + err_ba = F.normalize(layer[0], 127, 2) # [0, 255] + err_ab = F.normalize(layer[1], 127, 2) # [0, 255] + layer_data.append([err_ba, err_ab]) + + return layer_data + + +class Resize(object): + """Resize the input PIL Image to the given size. + + Args: + size (sequence or int): Desired output size. If size is a sequence like + (h, w), output size will be matched to this. If size is an int, + smaller edge of the image will be matched to this number. + i.e, if height > width, then image will be rescaled to + (size * height / width, size) + interpolation (int, optional): Desired interpolation. Default is + ``PIL.Image.BILINEAR`` + """ + + def __init__(self, size, interpolation=Image.BILINEAR): + assert isinstance(size, int) or (isinstance(size, collections.Iterable) and len(size) == 2) + self.size = size + self.interpolation = interpolation + + def __call__(self, inputs): + """ + Args: + img (PIL Image): Image to be scaled. + + Returns: + PIL Image: Rescaled image. + """ + return CustomFunc(inputs, F.resize, self.size, self.interpolation) + + +class RandomCrop(object): + """Crop the given PIL Image at a random location. + + Args: + size (sequence or int): Desired output size of the crop. If size is an + int instead of sequence like (h, w), a square crop (size, size) is + made. + padding (int or sequence, optional): Optional padding on each border + of the image. Default is 0, i.e no padding. If a sequence of length + 4 is provided, it is used to pad left, top, right, bottom borders + respectively. + """ + + def __init__(self, size, padding=0): + if isinstance(size, numbers.Number): + self.size = (int(size), int(size)) + else: + self.size = size + self.padding = padding + + @staticmethod + def get_params(img, output_size): + """Get parameters for ``crop`` for a random crop. + + Args: + img (PIL Image): Image to be cropped. + output_size (tuple): Expected output size of the crop. + + Returns: + tuple: params (i, j, h, w) to be passed to ``crop`` for random crop. + """ + w, h = img.size + th, tw = output_size + if w == tw and h == th: + return 0, 0, h, w + + i = random.randint(0, h - th) + j = random.randint(0, w - tw) + return i, j, th, tw + + def __call__(self, inputs): + """ + Args: + img (PIL Image): Image to be cropped. + + Returns: + PIL Image: Cropped image. + """ + if self.padding > 0: + inputs = CustomFunc(inputs, F.pad, self.padding) + + i, j, h, w = self.get_params(inputs[0], self.size) + return CustomFunc(inputs, F.crop, i, j, h, w) + + +class CenterCrop(object): + """Crop the given PIL Image at a random location. + + Args: + size (sequence or int): Desired output size of the crop. If size is an + int instead of sequence like (h, w), a square crop (size, size) is + made. + padding (int or sequence, optional): Optional padding on each border + of the image. Default is 0, i.e no padding. If a sequence of length + 4 is provided, it is used to pad left, top, right, bottom borders + respectively. + """ + + def __init__(self, size, padding=0): + if isinstance(size, numbers.Number): + self.size = (int(size), int(size)) + else: + self.size = size + self.padding = padding + + @staticmethod + def get_params(img, output_size): + """Get parameters for ``crop`` for a random crop. + + Args: + img (PIL Image): Image to be cropped. + output_size (tuple): Expected output size of the crop. + + Returns: + tuple: params (i, j, h, w) to be passed to ``crop`` for random crop. + """ + w, h = img.size + th, tw = output_size + if w == tw and h == th: + return 0, 0, h, w + + i = (h - th) // 2 + j = (w - tw) // 2 + return i, j, th, tw + + def __call__(self, inputs): + """ + Args: + img (PIL Image): Image to be cropped. + + Returns: + PIL Image: Cropped image. + """ + if self.padding > 0: + inputs = CustomFunc(inputs, F.pad, self.padding) + + i, j, h, w = self.get_params(inputs[0], self.size) + return CustomFunc(inputs, F.crop, i, j, h, w) + + +class RandomHorizontalFlip(object): + """Horizontally flip the given PIL Image randomly with a probability of 0.5.""" + + def __call__(self, inputs): + """ + Args: + img (PIL Image): Image to be flipped. + + Returns: + PIL Image: Randomly flipped image. + """ + + if random.random() < 0.5: + return CustomFunc(inputs, F.hflip) + return inputs + + +class RGB2Lab(object): + def __call__(self, inputs): + """ + Args: + img (PIL Image): Image to be flipped. + + Returns: + PIL Image: Randomly flipped image. + """ + + def __call__(self, inputs): + image_lab = color.rgb2lab(inputs[0]) + warp_ba_lab = color.rgb2lab(inputs[2]) + warp_aba_lab = color.rgb2lab(inputs[3]) + im_gbl_lab = color.rgb2lab(inputs[4]) + + inputs[0] = image_lab[:, :, :1] # l channel + inputs[1] = image_lab[:, :, 1:] # ab channel + inputs[2] = warp_ba_lab # lab channel + inputs[3] = warp_aba_lab # lab channel + inputs[4] = im_gbl_lab[:, :, 1:] # ab channel + + return inputs diff --git a/lib/videoloader.py b/lib/videoloader.py new file mode 100644 index 0000000..00c1d5d --- /dev/null +++ b/lib/videoloader.py @@ -0,0 +1,301 @@ +import sys + +sys.path.insert(0, "..") +import os +import random + +import cv2 +import numpy as np +import torch +import torchvision.utils as vutils +from PIL import Image +from skimage import color +from torch.autograd import Variable +from utils.flowlib import read_flow +from utils.util_distortion import CenterPad + +import lib.functional as F + +cv2.setNumThreads(0) + + +class RGB2Lab(object): + def __init__(self): + pass + + def __call__(self, inputs): + return color.rgb2lab(inputs) + + +class Normalize(object): + def __init__(self): + pass + + def __call__(self, inputs): + inputs[0:1, :, :] = F.normalize(inputs[0:1, :, :], 50, 1) + inputs[1:3, :, :] = F.normalize(inputs[1:3, :, :], (0, 0), (1, 1)) + return inputs + + +class ToTensor(object): + def __init__(self): + pass + + def __call__(self, inputs): + outputs = F.to_mytensor(inputs) # permute channel and transform to tensor + return outputs + + +class RandomErasing(object): + def __init__(self, probability=0.6, sl=0.05, sh=0.6): + self.probability = probability + self.sl = sl + self.sh = sh + + def __call__(self, img): + img = np.array(img) + if random.uniform(0, 1) > self.probability: + return Image.fromarray(img) + + area = img.shape[0] * img.shape[1] + h0 = img.shape[0] + w0 = img.shape[1] + channel = img.shape[2] + + h = int(round(random.uniform(self.sl, self.sh) * h0)) + w = int(round(random.uniform(self.sl, self.sh) * w0)) + + if w < img.shape[1] and h < img.shape[0]: + x1 = random.randint(0, img.shape[0] - h) + y1 = random.randint(0, img.shape[1] - w) + img[x1 : x1 + h, y1 : y1 + w, :] = np.random.rand(h, w, channel) * 255 + return Image.fromarray(img) + + return Image.fromarray(img) + + +class CenterCrop(object): + """ + center crop the numpy array + """ + + def __init__(self, image_size): + self.h0, self.w0 = image_size + + def __call__(self, input_numpy): + if input_numpy.ndim == 3: + h, w, channel = input_numpy.shape + output_numpy = np.zeros((self.h0, self.w0, channel)) + output_numpy = input_numpy[ + (h - self.h0) // 2 : (h - self.h0) // 2 + self.h0, (w - self.w0) // 2 : (w - self.w0) // 2 + self.w0, : + ] + else: + h, w = input_numpy.shape + output_numpy = np.zeros((self.h0, self.w0)) + output_numpy = input_numpy[ + (h - self.h0) // 2 : (h - self.h0) // 2 + self.h0, (w - self.w0) // 2 : (w - self.w0) // 2 + self.w0 + ] + return output_numpy + + +def parse_images(data_root): + image_pairs = [] + subdirs = sorted(os.listdir(data_root)) + for subdir in subdirs: + path = os.path.join(data_root, subdir) + if not os.path.isdir(path): + continue + + parse_file = os.path.join(path, "pairs_output_new.txt") + if os.path.exists(parse_file): + with open(parse_file, "r") as f: + lines = f.readlines() + for line in lines: + line = line.replace("\n", "") + ( + image1_name, + image2_name, + reference_video_name, + reference_video_name1, + reference_name1, + reference_name2, + reference_name3, + reference_name4, + reference_name5, + reference_gt1, + reference_gt2, + reference_gt3, + ) = line.split() + image1_name = image1_name.split(".")[0] + image2_name = image2_name.split(".")[0] + reference_video_name = reference_video_name.split(".")[0] + reference_video_name1 = reference_video_name1.split(".")[0] + reference_name1 = reference_name1.split(".")[0] + reference_name2 = reference_name2.split(".")[0] + reference_name3 = reference_name3.split(".")[0] + reference_name4 = reference_name4.split(".")[0] + reference_name5 = reference_name5.split(".")[0] + + reference_gt1 = reference_gt1.split(".")[0] + reference_gt2 = reference_gt2.split(".")[0] + reference_gt3 = reference_gt3.split(".")[0] + + flow_forward_name = image1_name + "_forward" + flow_backward_name = image1_name + "_backward" + mask_name = image1_name + "_mask" + + item = ( + image1_name + ".jpg", + image2_name + ".jpg", + reference_video_name + ".jpg", + reference_name1 + ".JPEG", + reference_name2 + ".JPEG", + reference_name3 + ".JPEG", + reference_name4 + ".JPEG", + reference_name5 + ".JPEG", + flow_forward_name + ".flo", + flow_backward_name + ".flo", + mask_name + ".pgm", + reference_gt1 + ".jpg", + reference_gt2 + ".jpg", + reference_gt3 + ".jpg", + path, + ) + image_pairs.append(item) + + else: + raise (RuntimeError("Error when parsing pair_output_count.txt in subfolders of: " + path + "\n")) + + return image_pairs + + +class VideosDataset(torch.utils.data.Dataset): + def __init__( + self, + data_root, + epoch, + image_size, + image_transform=None, + use_google_reference=False, + real_reference_probability=1, + nonzero_placeholder_probability=0.5, + ): + self.data_root = data_root + self.image_transform = image_transform + self.CenterPad = CenterPad(image_size) + self.ToTensor = ToTensor() + self.CenterCrop = CenterCrop(image_size) + + assert len(self.data_root) > 0, "find no dataroot" + self.epoch = epoch + self.image_pairs = parse_images(self.data_root) + self.real_len = len(self.image_pairs) + print("##### parsing image pairs in %s: %d pairs #####" % (data_root, self.real_len)) + self.image_pairs *= epoch + self.use_google_reference = use_google_reference + self.real_reference_probability = real_reference_probability + self.nonzero_placeholder_probability = nonzero_placeholder_probability + + def __getitem__(self, index): + ( + image1_name, + image2_name, + reference_video_name, + reference_name1, + reference_name2, + reference_name3, + reference_name4, + reference_name5, + flow_forward_name, + flow_backward_name, + mask_name, + reference_gt1, + reference_gt2, + reference_gt3, + path, + ) = self.image_pairs[index] + try: + I1 = Image.open(os.path.join(path, "input_pad", image1_name)) + I2 = Image.open(os.path.join(path, "input_pad", image2_name)) + + I_reference_video = Image.open( + os.path.join(path, "reference_gt", random.choice([reference_gt1, reference_gt2, reference_gt3])) + ) + I_reference_video_real = Image.open( + os.path.join( + path, + "reference", + random.choice( + [reference_name1, reference_name2, reference_name3, reference_name4, reference_name5] + ), + ) + ) + + flow_forward = read_flow(os.path.join(path, "flow", flow_forward_name)) # numpy + flow_backward = read_flow(os.path.join(path, "flow", flow_backward_name)) # numpy + mask = Image.open(os.path.join(path, "mask", mask_name)) + + # binary mask + mask = np.array(mask) + mask[mask < 240] = 0 + mask[mask >= 240] = 1 + + # transform + I1 = self.image_transform(I1) + I2 = self.image_transform(I2) + I_reference_video = self.image_transform(self.CenterPad(I_reference_video)) + I_reference_video_real = self.image_transform(self.CenterPad(I_reference_video_real)) + flow_forward = self.ToTensor(self.CenterCrop(flow_forward)) + flow_backward = self.ToTensor(self.CenterCrop(flow_backward)) + mask = self.ToTensor(self.CenterCrop(mask)) + + if np.random.random() < self.real_reference_probability: + I_reference_output = I_reference_video_real + placeholder = torch.zeros_like(I1) + self_ref_flag = torch.zeros_like(I1) + else: + I_reference_output = I_reference_video + placeholder = I2 if np.random.random() < self.nonzero_placeholder_probability else torch.zeros_like(I1) + self_ref_flag = torch.ones_like(I1) + + outputs = [ + I1, + I2, + I_reference_output, + flow_forward, + flow_backward, + mask, + placeholder, + self_ref_flag, + ] + + except Exception as e: + print("problem in, ", path) + print(e) + return self.__getitem__(np.random.randint(0, len(self.image_pairs))) + return outputs + + def __len__(self): + return len(self.image_pairs) + + +def batch_lab2rgb_transpose_mc(img_l_mc, img_ab_mc, nrow=8): + if isinstance(img_l_mc, Variable): + img_l_mc = img_l_mc.data.cpu() + if isinstance(img_ab_mc, Variable): + img_ab_mc = img_ab_mc.data.cpu() + + if img_l_mc.is_cuda: + img_l_mc = img_l_mc.cpu() + if img_ab_mc.is_cuda: + img_ab_mc = img_ab_mc.cpu() + + assert img_l_mc.dim() == 4 and img_ab_mc.dim() == 4, "only for batch input" + + l_norm, ab_norm = 1.0, 1.0 + l_mean, ab_mean = 50.0, 0 + img_l = img_l_mc * l_norm + l_mean + img_ab = img_ab_mc * ab_norm + ab_mean + pred_lab = torch.cat((img_l, img_ab), dim=1) + grid_lab = vutils.make_grid(pred_lab, nrow=nrow).numpy().astype("float64") + return (np.clip(color.lab2rgb(grid_lab.transpose((1, 2, 0))), 0, 1) * 255).astype("uint8") diff --git a/lib/videoloader_imagenet.py b/lib/videoloader_imagenet.py new file mode 100644 index 0000000..e7eb3c2 --- /dev/null +++ b/lib/videoloader_imagenet.py @@ -0,0 +1,290 @@ +import os +import os.path as osp +import struct + +import cv2 +import numpy as np +import torch +import torch.utils.data as data +import torchvision.transforms as transforms +from PIL import Image +from scipy.ndimage import map_coordinates +from scipy.ndimage.filters import gaussian_filter +from utils.util_distortion import CenterPadCrop_numpy, Distortion_with_flow, Normalize, RGB2Lab, ToTensor + +cv2.setNumThreads(0) + + +def parse_images(dir, with_bad, with_mid): + print("dir is: ", dir) + dir = osp.expanduser(dir) + + image_pairs = [] + + no_mid_cnt, no_bad_cnt = 0, 0 + for target in sorted(os.listdir(dir)): + d = osp.join(dir, target) + if not osp.isdir(d): + continue + + pair_file = osp.join(d, "pairs.txt") + + if osp.exists(pair_file): + with open(pair_file, "r") as f: + lines = f.readlines() + for line in lines: + pair = line.strip().split(" ") + name0, postfix = pair[0].split(".") + name1, postfix = pair[1].split(".") + if float(pair[2]) > 0: + item0 = (dir, target, name0, name1, 2) + item1 = (dir, target, name1, name0, 2) + image_pairs.append(item0) + image_pairs.append(item1) + + else: + raise (RuntimeError("Found no pair.txt in subfolders of: " + d + "\n")) + + if with_mid: + pair_file = osp.join(d, "pairs_mid.txt") + + if osp.exists(pair_file): + with open(pair_file, "r") as f: + for line in f: + pair = line.strip().split(" ") + name0, postfix = pair[0].split(".") + name1, postfix = pair[1].split(".") + item0 = (dir, target, name0, name1, 0) + item1 = (dir, target, name1, name0, 0) + image_pairs.append(item0) + image_pairs.append(item1) + + else: + no_mid_cnt += 1 + + if with_bad: + pair_file = osp.join(d, "pairs_bad.txt") + + if osp.exists(pair_file): + with open(pair_file, "r") as f: + for line in f: + pair = line.strip().split(" ") + name0, postfix = pair[0].split(".") + name1, postfix = pair[1].split(".") + item0 = (dir, target, name0, name1, -1) + item1 = (dir, target, name1, name0, -1) + image_pairs.append(item0) + image_pairs.append(item1) + + else: + no_bad_cnt += 1 + + if no_bad_cnt > 0: + print("find no pairs_bad.txt in %d folders for dir: %s" % (no_bad_cnt, dir)) + + if no_mid_cnt > 0: + print("find no pairs_mid.txt in %d folders for dir: %s" % (no_mid_cnt, dir)) + + return image_pairs + + +def pil_loader(path): + with open(path, "rb") as f: + with Image.open(f) as img: + return img.convert("RGB") + + +def image_loader(path): + return pil_loader(path) + + +def combo5_loader(path, real_w, real_h): + with open(path, "rb") as f: + d = f.read(4) + im_sz = struct.unpack("i", d) + h = im_sz[0] + + d = f.read(4) + im_sz = struct.unpack("i", d) + w = im_sz[0] + + d = f.read(4) + im_sz = struct.unpack("i", d) + d = f.read(im_sz[0]) + file_bytes = np.asarray(bytearray(d), dtype=np.uint8) + img_data_ndarray = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR) + img_data_ndarray = cv2.cvtColor(img_data_ndarray, cv2.COLOR_BGR2RGB) + warp_ba = Image.fromarray(img_data_ndarray) + img_data_ndarray = img_data_ndarray * 0 + warp_aba = Image.fromarray(img_data_ndarray) + + errs = [] + empty_array = np.zeros([np.array(warp_ba).shape[0], np.array(warp_ba).shape[1]]) + for _ in range(5): + err_ba = Image.fromarray(empty_array) + err_ab = Image.fromarray(empty_array) + errs.append([err_ba, err_ab]) + return errs, warp_ba, warp_aba + + +class VideosDataset_ImageNet(data.Dataset): + def __init__( + self, + data_root, + epoch, + image_size, + with_bad=False, + with_mid=False, + transforms_imagenet=None, + distortion_level=3, + brightnessjitter=0, + nonzero_placeholder_probability=0.5, + extra_reference_transform=None, + real_reference_probability=1, + ): + image_pairs = [] + curr_image_pairs = parse_images(data_root, with_bad, with_mid) + image_pairs += curr_image_pairs + print("##### parsing image_a pairs in %s: %d pairs #####" % (data_root, len(curr_image_pairs))) + if not image_pairs: + raise RuntimeError("Found 0 image_a pairs in all the data_roots") + + self.image_pairs = image_pairs + self.transforms_imagenet_raw = transforms_imagenet + self.extra_reference_transform = transforms.Compose(extra_reference_transform) + self.real_reference_probability = real_reference_probability + self.transforms_imagenet = transforms.Compose(transforms_imagenet) + self.epoch = epoch + self.image_size = image_size + self.real_len = len(self.image_pairs) + self.image_pairs *= epoch + self.distortion_level = distortion_level + self.distortion_transform = Distortion_with_flow() + self.brightnessjitter = brightnessjitter + self.flow_transform = transforms.Compose([CenterPadCrop_numpy(self.image_size), ToTensor()]) + self.nonzero_placeholder_probability = nonzero_placeholder_probability + self.ToTensor = ToTensor() + self.Normalize = Normalize() + + def __getitem__(self, index): + try: + image_id = 0 + pair_id = index + + combo_path = None + image_a_path = None + image_b_path = None + + image_names = ["", ""] + dir_root, cls_dir, image_names[0], image_names[1], is_good = self.image_pairs[pair_id] + sub_dir = osp.join(dir_root, cls_dir) + if is_good >= 1: + image_a_path = osp.join(sub_dir, "input", "%s.JPEG" % image_names[image_id]) + image_b_path = osp.join(sub_dir, "input", "%s.JPEG" % image_names[1 - image_id]) + elif is_good == 0: + image_a_path = osp.join(sub_dir, "input_mid", "%s.JPEG" % image_names[image_id]) + image_b_path = osp.join(sub_dir, "input_mid", "%s.JPEG" % image_names[1 - image_id]) + else: + image_a_path = osp.join(sub_dir, "input_bad", "%s.JPEG" % image_names[image_id]) + image_b_path = osp.join(sub_dir, "input_bad", "%s.JPEG" % image_names[1 - image_id]) + + if np.random.random() > 0.5: + image_a_path, image_b_path = image_b_path, image_a_path + + I1 = image_loader(image_a_path) + I2 = I1 + I_reference_video = I1 + I_reference_video_real = image_loader(image_b_path) + + ## generate the flow + height, width = np.array(I2).shape[0], np.array(I2).shape[1] + alpha = np.random.rand() * self.distortion_level + distortion_range = 50 + random_state = np.random.RandomState(None) + shape = self.image_size[0], self.image_size[1] + # dx: flow on the vertical direction; dy: flow on the horizontal direction + forward_dx = ( + gaussian_filter((random_state.rand(*shape) * 2 - 1), distortion_range, mode="constant", cval=0) + * alpha + * 1000 + ) + forward_dy = ( + gaussian_filter((random_state.rand(*shape) * 2 - 1), distortion_range, mode="constant", cval=0) + * alpha + * 1000 + ) + + for transform in self.transforms_imagenet_raw: + if type(transform) is RGB2Lab: + I1_raw = I1 + I1 = transform(I1) + for transform in self.transforms_imagenet_raw: + if type(transform) is RGB2Lab: + I2 = self.distortion_transform(I2, forward_dx, forward_dy) + I2_raw = I2 + I2 = transform(I2) + I2[0:1, :, :] = I2[0:1, :, :] + torch.randn(1) * self.brightnessjitter + + I_reference_video = self.extra_reference_transform(I_reference_video) + for transform in self.transforms_imagenet_raw: + I_reference_video = transform(I_reference_video) + + I_reference_video_real = self.transforms_imagenet(I_reference_video_real) + + flow_forward_raw = np.stack((forward_dy, forward_dx), axis=-1) + flow_backward_raw = np.zeros_like(flow_forward_raw) + flow_forward = self.flow_transform(flow_forward_raw) + flow_backward = self.flow_transform(flow_backward_raw) + + # update the mask for the pixels on the border + grid_x, grid_y = np.meshgrid(np.arange(self.image_size[0]), np.arange(self.image_size[1]), indexing="ij") + grid = np.stack((grid_y, grid_x), axis=-1) + grid_warp = grid + flow_forward_raw + location_y = grid_warp[:, :, 0].flatten() + location_x = grid_warp[:, :, 1].flatten() + I2_raw = np.array(I2_raw).astype(float) + I21_r = map_coordinates(I2_raw[:, :, 0], np.stack((location_x, location_y)), cval=-1).reshape( + (self.image_size[0], self.image_size[1]) + ) + I21_g = map_coordinates(I2_raw[:, :, 1], np.stack((location_x, location_y)), cval=-1).reshape( + (self.image_size[0], self.image_size[1]) + ) + I21_b = map_coordinates(I2_raw[:, :, 2], np.stack((location_x, location_y)), cval=-1).reshape( + (self.image_size[0], self.image_size[1]) + ) + I21_raw = np.stack((I21_r, I21_g, I21_b), axis=2) + mask = np.ones((self.image_size[0], self.image_size[1])) + mask[(I21_raw[:, :, 0] == -1) & (I21_raw[:, :, 1] == -1) & (I21_raw[:, :, 2] == -1)] = 0 + mask[abs(I21_raw - I1_raw).sum(axis=-1) > 50] = 0 + mask = self.ToTensor(mask) + + if np.random.random() < self.real_reference_probability: + I_reference_output = I_reference_video_real + placeholder = torch.zeros_like(I1) + self_ref_flag = torch.zeros_like(I1) + else: + I_reference_output = I_reference_video + placeholder = I2 if np.random.random() < self.nonzero_placeholder_probability else torch.zeros_like(I1) + self_ref_flag = torch.ones_like(I1) + + outputs = [ + I1, + I2, + I_reference_output, + flow_forward, + flow_backward, + mask, + placeholder, + self_ref_flag, + ] + + except Exception as e: + if combo_path is not None: + print("problem in ", combo_path) + print("problem in, ", image_a_path) + print(e) + return self.__getitem__(np.random.randint(0, len(self.image_pairs))) + return outputs + + def __len__(self): + return len(self.image_pairs) diff --git a/models/ContextualLoss.py b/models/ContextualLoss.py new file mode 100644 index 0000000..a37a8c6 --- /dev/null +++ b/models/ContextualLoss.py @@ -0,0 +1,303 @@ +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision +from utils.util import feature_normalize + +postpa = torchvision.transforms.Compose( + [ + torchvision.transforms.Lambda(lambda x: x.mul_(1.0 / 255)), + torchvision.transforms.Normalize( + mean=[-0.40760392, -0.45795686, -0.48501961], std=[1, 1, 1] # add imagenet mean + ), + torchvision.transforms.Lambda(lambda x: x[torch.LongTensor([2, 1, 0])]), # turn to RGB + ] +) +postpb = torchvision.transforms.Compose([torchvision.transforms.ToPILImage()]) + + +def post_processing(tensor): + t = postpa(tensor) # denormalize the image since the optimized tensor is the normalized one + t[t > 1] = 1 + t[t < 0] = 0 + img = postpb(t) + img = np.array(img) + return img + + +class ContextualLoss(nn.Module): + """ + input is Al, Bl, channel = 1, range ~ [0, 255] + """ + + def __init__(self): + super(ContextualLoss, self).__init__() + return None + + def forward(self, X_features, Y_features, h=0.1, feature_centering=True): + """ + X_features&Y_features are are feature vectors or feature 2d array + h: bandwidth + return the per-sample loss + """ + batch_size = X_features.shape[0] + feature_depth = X_features.shape[1] + feature_size = X_features.shape[2] + + # to normalized feature vectors + if feature_centering: + X_features = X_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze( + dim=-1 + ).unsqueeze(dim=-1) + Y_features = Y_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze( + dim=-1 + ).unsqueeze(dim=-1) + X_features = feature_normalize(X_features).view( + batch_size, feature_depth, -1 + ) # batch_size * feature_depth * feature_size^2 + Y_features = feature_normalize(Y_features).view( + batch_size, feature_depth, -1 + ) # batch_size * feature_depth * feature_size^2 + + # conine distance = 1 - similarity + X_features_permute = X_features.permute(0, 2, 1) # batch_size * feature_size^2 * feature_depth + d = 1 - torch.matmul(X_features_permute, Y_features) # batch_size * feature_size^2 * feature_size^2 + + # normalized distance: dij_bar + d_norm = d / (torch.min(d, dim=-1, keepdim=True)[0] + 1e-5) # batch_size * feature_size^2 * feature_size^2 + + # pairwise affinity + w = torch.exp((1 - d_norm) / h) + A_ij = w / torch.sum(w, dim=-1, keepdim=True) + + # contextual loss per sample + CX = torch.mean(torch.max(A_ij, dim=1)[0], dim=-1) + return -torch.log(CX) + + +class ContextualLoss_forward(nn.Module): + """ + input is Al, Bl, channel = 1, range ~ [0, 255] + """ + + def __init__(self): + super(ContextualLoss_forward, self).__init__() + return None + + def forward(self, X_features, Y_features, h=0.1, feature_centering=True): + """ + X_features&Y_features are are feature vectors or feature 2d array + h: bandwidth + return the per-sample loss + """ + batch_size = X_features.shape[0] + feature_depth = X_features.shape[1] + feature_size = X_features.shape[2] + + # to normalized feature vectors + if feature_centering: + X_features = X_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze( + dim=-1 + ).unsqueeze(dim=-1) + Y_features = Y_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze( + dim=-1 + ).unsqueeze(dim=-1) + X_features = feature_normalize(X_features).view( + batch_size, feature_depth, -1 + ) # batch_size * feature_depth * feature_size^2 + Y_features = feature_normalize(Y_features).view( + batch_size, feature_depth, -1 + ) # batch_size * feature_depth * feature_size^2 + + # conine distance = 1 - similarity + X_features_permute = X_features.permute(0, 2, 1) # batch_size * feature_size^2 * feature_depth + d = 1 - torch.matmul(X_features_permute, Y_features) # batch_size * feature_size^2 * feature_size^2 + + # normalized distance: dij_bar + d_norm = d / (torch.min(d, dim=-1, keepdim=True)[0] + 1e-5) # batch_size * feature_size^2 * feature_size^2 + + # pairwise affinity + w = torch.exp((1 - d_norm) / h) + A_ij = w / torch.sum(w, dim=-1, keepdim=True) + + # contextual loss per sample + CX = torch.mean(torch.max(A_ij, dim=-1)[0], dim=1) + return -torch.log(CX) + + +class ContextualLoss_complex(nn.Module): + """ + input is Al, Bl, channel = 1, range ~ [0, 255] + """ + + def __init__(self): + super(ContextualLoss_complex, self).__init__() + return None + + def forward(self, X_features, Y_features, h=0.1, patch_size=1, direction="forward"): + """ + X_features&Y_features are are feature vectors or feature 2d array + h: bandwidth + return the per-sample loss + """ + batch_size = X_features.shape[0] + feature_depth = X_features.shape[1] + feature_size = X_features.shape[2] + + # to normalized feature vectors + X_features = X_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze( + dim=-1 + ).unsqueeze(dim=-1) + Y_features = Y_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze( + dim=-1 + ).unsqueeze(dim=-1) + X_features = feature_normalize(X_features) # batch_size * feature_depth * feature_size^2 + Y_features = feature_normalize(Y_features) # batch_size * feature_depth * feature_size^2 + + # to normalized feature vectors + X_features = F.unfold( + X_features, kernel_size=(patch_size, patch_size), stride=(1, 1), padding=(patch_size // 2, patch_size // 2) + ) # batch_size * feature_depth_new * feature_size^2 + Y_features = F.unfold( + Y_features, kernel_size=(patch_size, patch_size), stride=(1, 1), padding=(patch_size // 2, patch_size // 2) + ) # batch_size * feature_depth_new * feature_size^2 + + # conine distance = 1 - similarity + X_features_permute = X_features.permute(0, 2, 1) # batch_size * feature_size^2 * feature_depth + d = 1 - torch.matmul(X_features_permute, Y_features) # batch_size * feature_size^2 * feature_size^2 + + # normalized distance: dij_bar + d_norm = d / (torch.min(d, dim=-1, keepdim=True)[0] + 1e-5) # batch_size * feature_size^2 * feature_size^2 + + # pairwise affinity + w = torch.exp((1 - d_norm) / h) + A_ij = w / torch.sum(w, dim=-1, keepdim=True) + + # contextual loss per sample + if direction == "forward": + CX = torch.mean(torch.max(A_ij, dim=-1)[0], dim=1) + else: + CX = torch.mean(torch.max(A_ij, dim=1)[0], dim=-1) + + return -torch.log(CX) + + +class ChamferDistance_patch_loss(nn.Module): + """ + input is Al, Bl, channel = 1, range ~ [0, 255] + """ + + def __init__(self): + super(ChamferDistance_patch_loss, self).__init__() + return None + + def forward(self, X_features, Y_features, patch_size=3, image_x=None, image_y=None, h=0.1, Y_features_in=None): + """ + X_features&Y_features are are feature vectors or feature 2d array + h: bandwidth + return the per-sample loss + """ + batch_size = X_features.shape[0] + feature_depth = X_features.shape[1] + feature_size = X_features.shape[2] + + # to normalized feature vectors + X_features = F.unfold( + X_features, kernel_size=(patch_size, patch_size), stride=(1, 1), padding=(patch_size // 2, patch_size // 2) + ) # batch_size, feature_depth_new * feature_size^2 + Y_features = F.unfold( + Y_features, kernel_size=(patch_size, patch_size), stride=(1, 1), padding=(patch_size // 2, patch_size // 2) + ) # batch_size, feature_depth_new * feature_size^2 + + if image_x is not None and image_y is not None: + image_x = torch.nn.functional.interpolate(image_x, size=(feature_size, feature_size), mode="bilinear").view( + batch_size, 3, -1 + ) + image_y = torch.nn.functional.interpolate(image_y, size=(feature_size, feature_size), mode="bilinear").view( + batch_size, 3, -1 + ) + + X_features_permute = X_features.permute(0, 2, 1) # batch_size * feature_size^2 * feature_depth + similarity_matrix = torch.matmul(X_features_permute, Y_features) # batch_size * feature_size^2 * feature_size^2 + NN_index = similarity_matrix.max(dim=-1, keepdim=True)[1].squeeze() + + if Y_features_in is not None: + loss = torch.mean((X_features - Y_features_in.detach()) ** 2) + Y_features_in = Y_features_in.detach() + else: + loss = torch.mean((X_features - Y_features[:, :, NN_index].detach()) ** 2) + Y_features_in = Y_features[:, :, NN_index].detach() + + # re-arrange image + if image_x is not None and image_y is not None: + image_y_rearrange = image_y[:, :, NN_index] + image_y_rearrange = image_y_rearrange.view(batch_size, 3, feature_size, feature_size) + image_x = image_x.view(batch_size, 3, feature_size, feature_size) + image_y = image_y.view(batch_size, 3, feature_size, feature_size) + + return loss + + +class ChamferDistance_loss(nn.Module): + """ + input is Al, Bl, channel = 1, range ~ [0, 255] + """ + + def __init__(self): + super(ChamferDistance_loss, self).__init__() + return None + + def forward(self, X_features, Y_features, image_x, image_y, h=0.1, Y_features_in=None): + """ + X_features&Y_features are are feature vectors or feature 2d array + h: bandwidth + return the per-sample loss + """ + batch_size = X_features.shape[0] + feature_depth = X_features.shape[1] + feature_size = X_features.shape[2] + + # to normalized feature vectors + X_features = feature_normalize(X_features).view( + batch_size, feature_depth, -1 + ) # batch_size * feature_depth * feature_size^2 + Y_features = feature_normalize(Y_features).view( + batch_size, feature_depth, -1 + ) # batch_size * feature_depth * feature_size^2 + image_x = torch.nn.functional.interpolate(image_x, size=(feature_size, feature_size), mode="bilinear").view( + batch_size, 3, -1 + ) + image_y = torch.nn.functional.interpolate(image_y, size=(feature_size, feature_size), mode="bilinear").view( + batch_size, 3, -1 + ) + + X_features_permute = X_features.permute(0, 2, 1) # batch_size * feature_size^2 * feature_depth + similarity_matrix = torch.matmul(X_features_permute, Y_features) # batch_size * feature_size^2 * feature_size^2 + NN_index = similarity_matrix.max(dim=-1, keepdim=True)[1].squeeze() + if Y_features_in is not None: + loss = torch.mean((X_features - Y_features_in.detach()) ** 2) + Y_features_in = Y_features_in.detach() + else: + loss = torch.mean((X_features - Y_features[:, :, NN_index].detach()) ** 2) + Y_features_in = Y_features[:, :, NN_index].detach() + + # re-arrange image + image_y_rearrange = image_y[:, :, NN_index] + image_y_rearrange = image_y_rearrange.view(batch_size, 3, feature_size, feature_size) + image_x = image_x.view(batch_size, 3, feature_size, feature_size) + image_y = image_y.view(batch_size, 3, feature_size, feature_size) + + return loss, Y_features_in, X_features + + +if __name__ == "__main__": + contextual_loss = ContextualLoss() + batch_size = 32 + feature_depth = 8 + feature_size = 16 + X_features = torch.zeros(batch_size, feature_depth, feature_size, feature_size) + Y_features = torch.zeros(batch_size, feature_depth, feature_size, feature_size) + + cx_loss = contextual_loss(X_features, Y_features, 1) + print(cx_loss) diff --git a/models/GAN_models.py b/models/GAN_models.py index a3b2f3d..3dddc07 100644 --- a/models/GAN_models.py +++ b/models/GAN_models.py @@ -101,61 +101,13 @@ class Self_Attention(nn.Module): return out -class Discriminator_WITHOUT_FC_x64_video(nn.Module): +class Discriminator_x64(nn.Module): """ Discriminative Network """ def __init__(self, in_size=6, ndf=64): - super(Discriminator_WITHOUT_FC_x64_video, self).__init__() - self.in_size = in_size - self.ndf = ndf - - self.layer1 = nn.Sequential( - # input size is in_size x 64 x 64 - SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)), - nn.LeakyReLU(0.2, inplace=True), - ) - self.layer2 = nn.Sequential( - # state size: ndf x 32 x 32 - SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf * 2), - nn.LeakyReLU(0.2, inplace=True), - ) - self.layer3 = nn.Sequential( - # state size: (ndf * 2) x 16 x 16 - SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf * 4), - nn.LeakyReLU(0.2, inplace=True), - ) - self.layer4 = nn.Sequential( - # state size: (ndf * 4) x 8 x 8 - # Self_Attention(self.ndf * 4, 'relu'), - SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf * 8), - nn.LeakyReLU(0.2, inplace=True), - ) - - self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, [3, 6], 1, 0)) - - def forward(self, input): - feature1 = self.layer1(input) - feature2 = self.layer2(feature1) - feature3 = self.layer3(feature2) - feature4 = self.layer4(feature3) - output = self.last(feature4) - output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1) - - return output, feature4 - - -class Discriminator_WITHOUT_FC_x64_video_BIG(nn.Module): - """ - Discriminative Network - """ - - def __init__(self, in_size=6, ndf=64): - super(Discriminator_WITHOUT_FC_x64_video_BIG, self).__init__() + super(Discriminator_x64, self).__init__() self.in_size = in_size self.ndf = ndf @@ -205,189 +157,237 @@ class Discriminator_WITHOUT_FC_x64_video_BIG(nn.Module): return output, feature4 -class Discriminator_WITHOUT_FC_x64(nn.Module): - """ - Discriminative Network - """ +# class Discriminator_WITHOUT_FC_x64_video(nn.Module): +# """ +# Discriminative Network +# """ - def __init__(self, in_size=3, ndf=64): - super(Discriminator_WITHOUT_FC_x64, self).__init__() - self.in_size = in_size - self.ndf = ndf +# def __init__(self, in_size=6, ndf=64): +# super(Discriminator_WITHOUT_FC_x64_video, self).__init__() +# self.in_size = in_size +# self.ndf = ndf - self.main = nn.Sequential( - # input size is in_size x 64 x 64 - SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)), - nn.LeakyReLU(0.2, inplace=True), - # state size: ndf x 32 x 32 - Self_Attention(self.ndf), - SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf * 2), - nn.LeakyReLU(0.2, inplace=True), - # state size: (ndf * 2) x 16 x 16 - SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf * 4), - nn.LeakyReLU(0.2, inplace=True), - # state size: (ndf * 4) x 8 x 8 - SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf * 8), - nn.LeakyReLU(0.2, inplace=True), - # sate size: (ndf * 8) x 4 x 4 - ) +# self.layer1 = nn.Sequential( +# # input size is in_size x 64 x 64 +# SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)), +# nn.LeakyReLU(0.2, inplace=True), +# ) +# self.layer2 = nn.Sequential( +# # state size: ndf x 32 x 32 +# SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf * 2), +# nn.LeakyReLU(0.2, inplace=True), +# ) +# self.layer3 = nn.Sequential( +# # state size: (ndf * 2) x 16 x 16 +# SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf * 4), +# nn.LeakyReLU(0.2, inplace=True), +# ) +# self.layer4 = nn.Sequential( +# # state size: (ndf * 4) x 8 x 8 +# # Self_Attention(self.ndf * 4, 'relu'), +# SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf * 8), +# nn.LeakyReLU(0.2, inplace=True), +# ) - self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0)) +# self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, [3, 6], 1, 0)) - def forward(self, input): +# def forward(self, input): +# feature1 = self.layer1(input) +# feature2 = self.layer2(feature1) +# feature3 = self.layer3(feature2) +# feature4 = self.layer4(feature3) +# output = self.last(feature4) +# output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1) - feature = self.main(input) - output = self.last(feature) - output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1) - - return output, feature +# return output, feature4 -class Discriminator_WITHOUT_FC_x64_BIG(nn.Module): - """ - Discriminative Network - """ +# class Discriminator_WITHOUT_FC_x64(nn.Module): +# """ +# Discriminative Network +# """ - def __init__(self, in_size=3, ndf=64): - super(Discriminator_WITHOUT_FC_x64_BIG, self).__init__() - self.in_size = in_size - self.ndf = ndf +# def __init__(self, in_size=3, ndf=64): +# super(Discriminator_WITHOUT_FC_x64, self).__init__() +# self.in_size = in_size +# self.ndf = ndf - self.main = nn.Sequential( - # 256 X 256 - SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)), - nn.LeakyReLU(0.2, inplace=True), - # 128 X 128 - SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf), - nn.LeakyReLU(0.2, inplace=True), - # 64 X 64 - Self_Attention(self.ndf), - SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf * 2), - nn.LeakyReLU(0.2, inplace=True), - # 32 X 32 - SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf * 4), - nn.LeakyReLU(0.2, inplace=True), - # 16 X 16 - SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf * 8), - nn.LeakyReLU(0.2, inplace=True), - # 8 X 8 - SpectralNorm(nn.Conv2d(self.ndf * 8, self.ndf * 16, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf * 16), - nn.LeakyReLU(0.2, inplace=True), - ) - # 4 X 4 - self.last = SpectralNorm(nn.Conv2d(self.ndf * 16, 1, 4, 1, 0)) +# self.main = nn.Sequential( +# # input size is in_size x 64 x 64 +# SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)), +# nn.LeakyReLU(0.2, inplace=True), +# # state size: ndf x 32 x 32 +# Self_Attention(self.ndf), +# SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf * 2), +# nn.LeakyReLU(0.2, inplace=True), +# # state size: (ndf * 2) x 16 x 16 +# SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf * 4), +# nn.LeakyReLU(0.2, inplace=True), +# # state size: (ndf * 4) x 8 x 8 +# SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf * 8), +# nn.LeakyReLU(0.2, inplace=True), +# # sate size: (ndf * 8) x 4 x 4 +# ) - def forward(self, input): +# self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0)) - feature = self.main(input) - output = self.last(feature) - output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1) +# def forward(self, input): - return output, feature +# feature = self.main(input) +# output = self.last(feature) +# output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1) + +# return output, feature -class Discriminator_WITHOUT_FC_x128(nn.Module): - """ - Discriminative Network - """ +# class Discriminator_WITHOUT_FC_x64_BIG(nn.Module): +# """ +# Discriminative Network +# """ - def __init__(self, in_size=3, ndf=64): - super(Discriminator_WITHOUT_FC_x128, self).__init__() - self.in_size = in_size - self.ndf = ndf +# def __init__(self, in_size=3, ndf=64): +# super(Discriminator_WITHOUT_FC_x64_BIG, self).__init__() +# self.in_size = in_size +# self.ndf = ndf - self.main = nn.Sequential( - # input size is in_size x 128 x 128 - SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)), - nn.LeakyReLU(0.2, inplace=True), - # state size: ndf x 64 x 64 - SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf), - nn.LeakyReLU(0.2, inplace=True), - # state size: ndf x 32 x 32 - SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf * 2), - nn.LeakyReLU(0.2, inplace=True), - # state size: (ndf * 24) x 16 x 16 - SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf * 4), - nn.LeakyReLU(0.2, inplace=True), - # sate size: (ndf * 4) x 8 x 8 - SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf * 8), - nn.LeakyReLU(0.2, inplace=True), - # sate size: (ndf * 8) x 4 x 4 - ) +# self.main = nn.Sequential( +# # 256 X 256 +# SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)), +# nn.LeakyReLU(0.2, inplace=True), +# # 128 X 128 +# SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf), +# nn.LeakyReLU(0.2, inplace=True), +# # 64 X 64 +# Self_Attention(self.ndf), +# SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf * 2), +# nn.LeakyReLU(0.2, inplace=True), +# # 32 X 32 +# SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf * 4), +# nn.LeakyReLU(0.2, inplace=True), +# # 16 X 16 +# SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf * 8), +# nn.LeakyReLU(0.2, inplace=True), +# # 8 X 8 +# SpectralNorm(nn.Conv2d(self.ndf * 8, self.ndf * 16, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf * 16), +# nn.LeakyReLU(0.2, inplace=True), +# ) +# # 4 X 4 +# self.last = SpectralNorm(nn.Conv2d(self.ndf * 16, 1, 4, 1, 0)) - self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0)) +# def forward(self, input): - def forward(self, input): +# feature = self.main(input) +# output = self.last(feature) +# output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1) - feature = self.main(input) - output = self.last(feature) - - return output, feature +# return output, feature -class Discriminator_WITHOUT_FC_x256(nn.Module): - """ - Discriminative Network - """ +# class Discriminator_WITHOUT_FC_x128(nn.Module): +# """ +# Discriminative Network +# """ - def __init__(self, in_size=3, ndf=64): - super(Discriminator_WITHOUT_FC_x256, self).__init__() - self.in_size = in_size - self.ndf = ndf +# def __init__(self, in_size=3, ndf=64): +# super(Discriminator_WITHOUT_FC_x128, self).__init__() +# self.in_size = in_size +# self.ndf = ndf - self.main = nn.Sequential( - # input size is in_size x 256 x 256 - SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)), - nn.LeakyReLU(0.2, inplace=True), - # state size: ndf x 128 x 128 - SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf), - nn.LeakyReLU(0.2, inplace=True), - # state size: ndf x 64 x 64 - SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf), - nn.LeakyReLU(0.2, inplace=True), - # state size: (ndf * 1) x 32 x 32 - SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf * 2), - nn.LeakyReLU(0.2, inplace=True), - # state size: (ndf * 2) x 16 x 16 - SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf * 4), - nn.LeakyReLU(0.2, inplace=True), - # sate size: (ndf * 4) x 8 x 8 - SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)), - nn.InstanceNorm2d(self.ndf * 8), - nn.LeakyReLU(0.2, inplace=True), - # sate size: (ndf * 8) x 4 x 4 - ) +# self.main = nn.Sequential( +# # input size is in_size x 128 x 128 +# SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)), +# nn.LeakyReLU(0.2, inplace=True), +# # state size: ndf x 64 x 64 +# SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf), +# nn.LeakyReLU(0.2, inplace=True), +# # state size: ndf x 32 x 32 +# SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf * 2), +# nn.LeakyReLU(0.2, inplace=True), +# # state size: (ndf * 24) x 16 x 16 +# SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf * 4), +# nn.LeakyReLU(0.2, inplace=True), +# # sate size: (ndf * 4) x 8 x 8 +# SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf * 8), +# nn.LeakyReLU(0.2, inplace=True), +# # sate size: (ndf * 8) x 4 x 4 +# ) - self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0)) +# self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0)) - def forward(self, input): +# def forward(self, input): - feature = self.main(input) - output = self.last(feature) +# feature = self.main(input) +# output = self.last(feature) - return output, feature +# return output, feature -if __name__ == "__main__": - discriminator1 = Discriminator_WITHOUT_FC_x64() - discriminator2 = Discriminator_WITHOUT_FC_x128() - discriminator3 = Discriminator_WITHOUT_FC_x256() - print(discriminator1) - print(discriminator2) - print(discriminator3) +# class Discriminator_WITHOUT_FC_x256(nn.Module): +# """ +# Discriminative Network +# """ + +# def __init__(self, in_size=3, ndf=64): +# super(Discriminator_WITHOUT_FC_x256, self).__init__() +# self.in_size = in_size +# self.ndf = ndf + +# self.main = nn.Sequential( +# # input size is in_size x 256 x 256 +# SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)), +# nn.LeakyReLU(0.2, inplace=True), +# # state size: ndf x 128 x 128 +# SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf), +# nn.LeakyReLU(0.2, inplace=True), +# # state size: ndf x 64 x 64 +# SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf), +# nn.LeakyReLU(0.2, inplace=True), +# # state size: (ndf * 1) x 32 x 32 +# SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf * 2), +# nn.LeakyReLU(0.2, inplace=True), +# # state size: (ndf * 2) x 16 x 16 +# SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf * 4), +# nn.LeakyReLU(0.2, inplace=True), +# # sate size: (ndf * 4) x 8 x 8 +# SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)), +# nn.InstanceNorm2d(self.ndf * 8), +# nn.LeakyReLU(0.2, inplace=True), +# # sate size: (ndf * 8) x 4 x 4 +# ) + +# self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0)) + +# def forward(self, input): + +# feature = self.main(input) +# output = self.last(feature) + +# return output, feature + + +# if __name__ == "__main__": +# discriminator1 = Discriminator_WITHOUT_FC_x64() +# discriminator2 = Discriminator_WITHOUT_FC_x128() +# discriminator3 = Discriminator_WITHOUT_FC_x256() +# print(discriminator1) +# print(discriminator2) +# print(discriminator3) diff --git a/models/NonlocalNet.py b/models/NonlocalNet.py index 1713db0..6683d1a 100644 --- a/models/NonlocalNet.py +++ b/models/NonlocalNet.py @@ -189,13 +189,13 @@ class Self_Attn(nn.Module): return out -class VGG19_feature_color_torchversion(nn.Module): +class VGG19_pytorch(nn.Module): """ NOTE: no need to pre-process the input; input tensor should range in [0,1] """ def __init__(self, pool="max"): - super(VGG19_feature_color_torchversion, self).__init__() + super(VGG19_pytorch, self).__init__() self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, padding=1) self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, padding=1) self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, padding=1) diff --git a/opt.txt b/opt.txt new file mode 100644 index 0000000..3d980bc --- /dev/null +++ b/opt.txt @@ -0,0 +1,37 @@ +train.py 2020-12-21 15:03:17 +------------ Options ------------- +batch_size: 16 +beta1: 0.5 +checkpoint_dir: checkpoints/default +contextual_loss_direction: forward +data_root: /home/zhanbo/remote/video/video_pair3/ +data_root_imagenet: /home/zhanbo/remote/imagenet_ref_pair/ +domain_invariant: False +epoch: 40 +gpu_ids: 0,1,2,3 +ic: 7 +image_size: [216, 384] +load_pretrained_model: True +lr: 1e-05 +lr_gamma: 0.1 +lr_step: 100 +luminance_noise: 2.0 +nonzero_placeholder_probability: 0.0 +permute_data: True +print_step: 2 +real_reference_probability: 0.7 +resume: False +resume_epoch: 0 +tb_log_step: 50 +weight_consistent: 0.02 +weight_contextual: 0.2 +weight_gan: 0.2 +weight_nonlocal_consistent: 0.0 +weight_nonlocal_smoothness: 0.0 +weight_perceptual: 0.001 +weight_smoothness: 5.0 +weigth_l1: 2.0 +with_bad: True +with_mid: True +workers: 8 +-------------- End ---------------- diff --git a/requirements.txt b/requirements.txt index 7539545..b940bde 100644 --- a/requirements.txt +++ b/requirements.txt @@ -14,4 +14,5 @@ prefetch_generator PyYAML tensorboard future -numba \ No newline at end of file +numba +pypng \ No newline at end of file diff --git a/sample_videos/output/v32_01/00000.jpg b/sample_videos/output/v32_01/00000.jpg new file mode 100644 index 0000000..49d9a3d Binary files /dev/null and b/sample_videos/output/v32_01/00000.jpg differ diff --git a/sample_videos/output/v32_01/00001.jpg b/sample_videos/output/v32_01/00001.jpg new file mode 100644 index 0000000..9bdfff7 Binary files /dev/null and b/sample_videos/output/v32_01/00001.jpg differ diff --git a/sample_videos/output/v32_01/00002.jpg b/sample_videos/output/v32_01/00002.jpg new file mode 100644 index 0000000..0c388a2 Binary files /dev/null and b/sample_videos/output/v32_01/00002.jpg differ diff --git a/sample_videos/output/v32_01/00003.jpg b/sample_videos/output/v32_01/00003.jpg new file mode 100644 index 0000000..8eb6a49 Binary files /dev/null and b/sample_videos/output/v32_01/00003.jpg differ diff --git a/sample_videos/output/v32_01/00004.jpg b/sample_videos/output/v32_01/00004.jpg new file mode 100644 index 0000000..7607dc1 Binary files /dev/null and b/sample_videos/output/v32_01/00004.jpg differ diff --git a/sample_videos/output/v32_01/00005.jpg b/sample_videos/output/v32_01/00005.jpg new file mode 100644 index 0000000..21cc55e Binary files /dev/null and b/sample_videos/output/v32_01/00005.jpg differ diff --git a/sample_videos/output/v32_01/00006.jpg b/sample_videos/output/v32_01/00006.jpg new file mode 100644 index 0000000..8b5ad49 Binary files /dev/null and b/sample_videos/output/v32_01/00006.jpg differ diff --git a/sample_videos/output/v32_01/00007.jpg b/sample_videos/output/v32_01/00007.jpg new file mode 100644 index 0000000..2c4fa2d Binary files /dev/null and b/sample_videos/output/v32_01/00007.jpg differ diff --git a/sample_videos/output/v32_01/00008.jpg b/sample_videos/output/v32_01/00008.jpg new file mode 100644 index 0000000..d1bb6e7 Binary files /dev/null and b/sample_videos/output/v32_01/00008.jpg differ diff --git a/sample_videos/output/v32_01/00009.jpg b/sample_videos/output/v32_01/00009.jpg new file mode 100644 index 0000000..17d2718 Binary files /dev/null and b/sample_videos/output/v32_01/00009.jpg differ diff --git a/sample_videos/output/v32_01/00010.jpg b/sample_videos/output/v32_01/00010.jpg new file mode 100644 index 0000000..2c7583e Binary files /dev/null and b/sample_videos/output/v32_01/00010.jpg differ diff --git a/sample_videos/output/v32_01/00011.jpg b/sample_videos/output/v32_01/00011.jpg new file mode 100644 index 0000000..95268e9 Binary files /dev/null and b/sample_videos/output/v32_01/00011.jpg differ diff --git a/sample_videos/output/v32_01/00012.jpg b/sample_videos/output/v32_01/00012.jpg new file mode 100644 index 0000000..88cbb6d Binary files /dev/null and b/sample_videos/output/v32_01/00012.jpg differ diff --git a/sample_videos/output/v32_01/00013.jpg b/sample_videos/output/v32_01/00013.jpg new file mode 100644 index 0000000..79353fd Binary files /dev/null and b/sample_videos/output/v32_01/00013.jpg differ diff --git a/tensorboardX/__init__.py b/tensorboardX/__init__.py new file mode 100644 index 0000000..e910378 --- /dev/null +++ b/tensorboardX/__init__.py @@ -0,0 +1,5 @@ +"""A module for visualization with tensorboard +""" + +from .writer import FileWriter, SummaryWriter +from .record_writer import RecordWriter diff --git a/tensorboardX/crc32c.py b/tensorboardX/crc32c.py new file mode 100644 index 0000000..62c9e0c --- /dev/null +++ b/tensorboardX/crc32c.py @@ -0,0 +1,123 @@ +import array + + +CRC_TABLE = ( + 0x00000000, 0xf26b8303, 0xe13b70f7, 0x1350f3f4, + 0xc79a971f, 0x35f1141c, 0x26a1e7e8, 0xd4ca64eb, + 0x8ad958cf, 0x78b2dbcc, 0x6be22838, 0x9989ab3b, + 0x4d43cfd0, 0xbf284cd3, 0xac78bf27, 0x5e133c24, + 0x105ec76f, 0xe235446c, 0xf165b798, 0x030e349b, + 0xd7c45070, 0x25afd373, 0x36ff2087, 0xc494a384, + 0x9a879fa0, 0x68ec1ca3, 0x7bbcef57, 0x89d76c54, + 0x5d1d08bf, 0xaf768bbc, 0xbc267848, 0x4e4dfb4b, + 0x20bd8ede, 0xd2d60ddd, 0xc186fe29, 0x33ed7d2a, + 0xe72719c1, 0x154c9ac2, 0x061c6936, 0xf477ea35, + 0xaa64d611, 0x580f5512, 0x4b5fa6e6, 0xb93425e5, + 0x6dfe410e, 0x9f95c20d, 0x8cc531f9, 0x7eaeb2fa, + 0x30e349b1, 0xc288cab2, 0xd1d83946, 0x23b3ba45, + 0xf779deae, 0x05125dad, 0x1642ae59, 0xe4292d5a, + 0xba3a117e, 0x4851927d, 0x5b016189, 0xa96ae28a, + 0x7da08661, 0x8fcb0562, 0x9c9bf696, 0x6ef07595, + 0x417b1dbc, 0xb3109ebf, 0xa0406d4b, 0x522bee48, + 0x86e18aa3, 0x748a09a0, 0x67dafa54, 0x95b17957, + 0xcba24573, 0x39c9c670, 0x2a993584, 0xd8f2b687, + 0x0c38d26c, 0xfe53516f, 0xed03a29b, 0x1f682198, + 0x5125dad3, 0xa34e59d0, 0xb01eaa24, 0x42752927, + 0x96bf4dcc, 0x64d4cecf, 0x77843d3b, 0x85efbe38, + 0xdbfc821c, 0x2997011f, 0x3ac7f2eb, 0xc8ac71e8, + 0x1c661503, 0xee0d9600, 0xfd5d65f4, 0x0f36e6f7, + 0x61c69362, 0x93ad1061, 0x80fde395, 0x72966096, + 0xa65c047d, 0x5437877e, 0x4767748a, 0xb50cf789, + 0xeb1fcbad, 0x197448ae, 0x0a24bb5a, 0xf84f3859, + 0x2c855cb2, 0xdeeedfb1, 0xcdbe2c45, 0x3fd5af46, + 0x7198540d, 0x83f3d70e, 0x90a324fa, 0x62c8a7f9, + 0xb602c312, 0x44694011, 0x5739b3e5, 0xa55230e6, + 0xfb410cc2, 0x092a8fc1, 0x1a7a7c35, 0xe811ff36, + 0x3cdb9bdd, 0xceb018de, 0xdde0eb2a, 0x2f8b6829, + 0x82f63b78, 0x709db87b, 0x63cd4b8f, 0x91a6c88c, + 0x456cac67, 0xb7072f64, 0xa457dc90, 0x563c5f93, + 0x082f63b7, 0xfa44e0b4, 0xe9141340, 0x1b7f9043, + 0xcfb5f4a8, 0x3dde77ab, 0x2e8e845f, 0xdce5075c, + 0x92a8fc17, 0x60c37f14, 0x73938ce0, 0x81f80fe3, + 0x55326b08, 0xa759e80b, 0xb4091bff, 0x466298fc, + 0x1871a4d8, 0xea1a27db, 0xf94ad42f, 0x0b21572c, + 0xdfeb33c7, 0x2d80b0c4, 0x3ed04330, 0xccbbc033, + 0xa24bb5a6, 0x502036a5, 0x4370c551, 0xb11b4652, + 0x65d122b9, 0x97baa1ba, 0x84ea524e, 0x7681d14d, + 0x2892ed69, 0xdaf96e6a, 0xc9a99d9e, 0x3bc21e9d, + 0xef087a76, 0x1d63f975, 0x0e330a81, 0xfc588982, + 0xb21572c9, 0x407ef1ca, 0x532e023e, 0xa145813d, + 0x758fe5d6, 0x87e466d5, 0x94b49521, 0x66df1622, + 0x38cc2a06, 0xcaa7a905, 0xd9f75af1, 0x2b9cd9f2, + 0xff56bd19, 0x0d3d3e1a, 0x1e6dcdee, 0xec064eed, + 0xc38d26c4, 0x31e6a5c7, 0x22b65633, 0xd0ddd530, + 0x0417b1db, 0xf67c32d8, 0xe52cc12c, 0x1747422f, + 0x49547e0b, 0xbb3ffd08, 0xa86f0efc, 0x5a048dff, + 0x8ecee914, 0x7ca56a17, 0x6ff599e3, 0x9d9e1ae0, + 0xd3d3e1ab, 0x21b862a8, 0x32e8915c, 0xc083125f, + 0x144976b4, 0xe622f5b7, 0xf5720643, 0x07198540, + 0x590ab964, 0xab613a67, 0xb831c993, 0x4a5a4a90, + 0x9e902e7b, 0x6cfbad78, 0x7fab5e8c, 0x8dc0dd8f, + 0xe330a81a, 0x115b2b19, 0x020bd8ed, 0xf0605bee, + 0x24aa3f05, 0xd6c1bc06, 0xc5914ff2, 0x37faccf1, + 0x69e9f0d5, 0x9b8273d6, 0x88d28022, 0x7ab90321, + 0xae7367ca, 0x5c18e4c9, 0x4f48173d, 0xbd23943e, + 0xf36e6f75, 0x0105ec76, 0x12551f82, 0xe03e9c81, + 0x34f4f86a, 0xc69f7b69, 0xd5cf889d, 0x27a40b9e, + 0x79b737ba, 0x8bdcb4b9, 0x988c474d, 0x6ae7c44e, + 0xbe2da0a5, 0x4c4623a6, 0x5f16d052, 0xad7d5351, +) + + +CRC_INIT = 0 + +_MASK = 0xFFFFFFFF + + +def crc_update(crc, data): + """Update CRC-32C checksum with data. + + Args: + crc: 32-bit checksum to update as long. + data: byte array, string or iterable over bytes. + + Returns: + 32-bit updated CRC-32C as long. + """ + + if type(data) != array.array or data.itemsize != 1: + buf = array.array("B", data) + else: + buf = data + + crc ^= _MASK + for b in buf: + table_index = (crc ^ b) & 0xff + crc = (CRC_TABLE[table_index] ^ (crc >> 8)) & _MASK + return crc ^ _MASK + + +def crc_finalize(crc): + """Finalize CRC-32C checksum. + + This function should be called as last step of crc calculation. + + Args: + crc: 32-bit checksum as long. + + Returns: + finalized 32-bit checksum as long + """ + return crc & _MASK + + +def crc32c(data): + """Compute CRC-32C checksum of the data. + + Args: + data: byte array, string or iterable over bytes. + + Returns: + 32-bit CRC-32C checksum of data as long. + """ + return crc_finalize(crc_update(CRC_INIT, data)) diff --git a/tensorboardX/embedding.py b/tensorboardX/embedding.py new file mode 100644 index 0000000..4e1935f --- /dev/null +++ b/tensorboardX/embedding.py @@ -0,0 +1,49 @@ +import os + + +def make_tsv(metadata, save_path): + metadata = [str(x) for x in metadata] + with open(os.path.join(save_path, 'metadata.tsv'), 'w') as f: + for x in metadata: + f.write(x + '\n') + + +# https://github.com/tensorflow/tensorboard/issues/44 image label will be squared +def make_sprite(label_img, save_path): + import math + import torch + import torchvision + from .x2num import makenp + # this ensures the sprite image has correct dimension as described in + # https://www.tensorflow.org/get_started/embedding_viz + nrow = int(math.ceil((label_img.size(0)) ** 0.5)) + + label_img = torch.from_numpy(makenp(label_img)) # for other framework + # augment images so that #images equals nrow*nrow + label_img = torch.cat((label_img, torch.randn(nrow ** 2 - label_img.size(0), *label_img.size()[1:]) * 255), 0) + + torchvision.utils.save_image(label_img, os.path.join(save_path, 'sprite.png'), nrow=nrow, padding=0) + + +def append_pbtxt(metadata, label_img, save_path, global_step, tag): + with open(os.path.join(save_path, 'projector_config.pbtxt'), 'a') as f: + # step = os.path.split(save_path)[-1] + f.write('embeddings {\n') + f.write('tensor_name: "{}:{}"\n'.format(tag, global_step)) + f.write('tensor_path: "{}"\n'.format(os.path.join(global_step, 'tensors.tsv'))) + if metadata is not None: + f.write('metadata_path: "{}"\n'.format(os.path.join(global_step, 'metadata.tsv'))) + if label_img is not None: + f.write('sprite {\n') + f.write('image_path: "{}"\n'.format(os.path.join(global_step, 'sprite.png'))) + f.write('single_image_dim: {}\n'.format(label_img.size(3))) + f.write('single_image_dim: {}\n'.format(label_img.size(2))) + f.write('}\n') + f.write('}\n') + + +def make_mat(matlist, save_path): + with open(os.path.join(save_path, 'tensors.tsv'), 'w') as f: + for x in matlist: + x = [str(i) for i in x] + f.write('\t'.join(x) + '\n') diff --git a/tensorboardX/event_file_writer.py b/tensorboardX/event_file_writer.py new file mode 100644 index 0000000..0f99838 --- /dev/null +++ b/tensorboardX/event_file_writer.py @@ -0,0 +1,196 @@ +# Copyright 2015 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Writes events to disk in a logdir.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import logging +import os.path +import socket +import threading +import time + +import six + +from .src import event_pb2 +from .record_writer import RecordWriter + + +def directory_check(path): + '''Initialize the directory for log files.''' + # If the direcotry does not exist, create it! + if not os.path.exists(path): + os.makedirs(path) + + +class EventsWriter(object): + '''Writes `Event` protocol buffers to an event file.''' + + def __init__(self, file_prefix): + ''' + Events files have a name of the form + '/some/file/path/events.out.tfevents.[timestamp].[hostname]' + ''' + self._file_prefix = file_prefix + ".out.tfevents." + str(time.time())[:10] + "." + socket.gethostname() + + # Open(Create) the log file with the particular form of name. + logging.basicConfig(filename=self._file_prefix) + + self._num_outstanding_events = 0 + + self._py_recordio_writer = RecordWriter(self._file_prefix) + + # Initialize an event instance. + self._event = event_pb2.Event() + + self._event.wall_time = time.time() + + self.write_event(self._event) + + def write_event(self, event): + '''Append "event" to the file.''' + + # Check if event is of type event_pb2.Event proto. + if not isinstance(event, event_pb2.Event): + raise TypeError("Expected an event_pb2.Event proto, " + " but got %s" % type(event)) + return self._write_serialized_event(event.SerializeToString()) + + def _write_serialized_event(self, event_str): + self._num_outstanding_events += 1 + self._py_recordio_writer.write(event_str) + + def flush(self): + '''Flushes the event file to disk.''' + self._num_outstanding_events = 0 + return True + + def close(self): + '''Call self.flush().''' + return_value = self.flush() + return return_value + + +class EventFileWriter(object): + """Writes `Event` protocol buffers to an event file. + The `EventFileWriter` class creates an event file in the specified directory, + and asynchronously writes Event protocol buffers to the file. The Event file + is encoded using the tfrecord format, which is similar to RecordIO. + @@__init__ + @@add_event + @@flush + @@close + """ + + def __init__(self, logdir, max_queue=10, flush_secs=120): + """Creates a `EventFileWriter` and an event file to write to. + On construction the summary writer creates a new event file in `logdir`. + This event file will contain `Event` protocol buffers, which are written to + disk via the add_event method. + The other arguments to the constructor control the asynchronous writes to + the event file: + * `flush_secs`: How often, in seconds, to flush the added summaries + and events to disk. + * `max_queue`: Maximum number of summaries or events pending to be + written to disk before one of the 'add' calls block. + Args: + logdir: A string. Directory where event file will be written. + max_queue: Integer. Size of the queue for pending events and summaries. + flush_secs: Number. How often, in seconds, to flush the + pending events and summaries to disk. + """ + self._logdir = logdir + directory_check(self._logdir) + self._event_queue = six.moves.queue.Queue(max_queue) + self._ev_writer = EventsWriter(os.path.join(self._logdir, "events")) + self._closed = False + self._worker = _EventLoggerThread(self._event_queue, self._ev_writer, + flush_secs) + + self._worker.start() + + def get_logdir(self): + """Returns the directory where event file will be written.""" + return self._logdir + + def reopen(self): + """Reopens the EventFileWriter. + Can be called after `close()` to add more events in the same directory. + The events will go into a new events file. + Does nothing if the EventFileWriter was not closed. + """ + if self._closed: + self._closed = False + + def add_event(self, event): + """Adds an event to the event file. + Args: + event: An `Event` protocol buffer. + """ + if not self._closed: + self._event_queue.put(event) + + def flush(self): + """Flushes the event file to disk. + Call this method to make sure that all pending events have been written to + disk. + """ + self._event_queue.join() + self._ev_writer.flush() + + def close(self): + """Flushes the event file to disk and close the file. + Call this method when you do not need the summary writer anymore. + """ + self.flush() + self._ev_writer.close() + self._closed = True + + +class _EventLoggerThread(threading.Thread): + """Thread that logs events.""" + + def __init__(self, queue, ev_writer, flush_secs): + """Creates an _EventLoggerThread. + Args: + queue: A Queue from which to dequeue events. + ev_writer: An event writer. Used to log brain events for + the visualizer. + flush_secs: How often, in seconds, to flush the + pending file to disk. + """ + threading.Thread.__init__(self) + self.daemon = True + self._queue = queue + self._ev_writer = ev_writer + self._flush_secs = flush_secs + # The first event will be flushed immediately. + self._next_event_flush_time = 0 + + def run(self): + while True: + event = self._queue.get() + try: + self._ev_writer.write_event(event) + # Flush the event writer every so often. + now = time.time() + if now > self._next_event_flush_time: + self._ev_writer.flush() + # Do it again in two minutes. + self._next_event_flush_time = now + self._flush_secs + finally: + self._queue.task_done() diff --git a/tensorboardX/graph.py b/tensorboardX/graph.py new file mode 100644 index 0000000..8135fd3 --- /dev/null +++ b/tensorboardX/graph.py @@ -0,0 +1,66 @@ +from .src.graph_pb2 import GraphDef +from .src.node_def_pb2 import NodeDef +from .src.versions_pb2 import VersionDef +from .src.attr_value_pb2 import AttrValue +from .src.tensor_shape_pb2 import TensorShapeProto + +from distutils.version import LooseVersion + + +def replace(name, scope): + return '/'.join([scope[name], name]) + + +def parse(graph): + import torch + scope = {} + for n in graph.nodes(): + inputs = [i.uniqueName() for i in n.inputs()] + for i in range(1, len(inputs)): + scope[inputs[i]] = n.scopeName() + + uname = next(iter(n.outputs())).uniqueName() + assert n.scopeName() != '', '{} has empty scope name'.format(n) + scope[uname] = n.scopeName() + if LooseVersion(torch.__version__) >= LooseVersion("0.4"): + scope['0'] = 'input' + else: + scope['1'] = 'input' + + nodes = [] + for n in graph.nodes(): + attrs = {k: n[k] for k in n.attributeNames()} + attrs = str(attrs).replace("'", ' ') # singlequote will be escaped by tensorboard + if any(i.uniqueName() not in scope.keys() for i in n.inputs()): # 0.3.1 workaround + continue + inputs = [replace(i.uniqueName(), scope) for i in n.inputs()] + uname = next(iter(n.outputs())).uniqueName() # FIXME: only first output is considered + nodes.append({'name': replace(uname, scope), 'op': n.kind(), 'inputs': inputs, 'attr': attrs}) + + for n in graph.inputs(): + uname = n.uniqueName() + if uname not in scope.keys(): + scope[uname] = 'unused' + nodes.append({'name': replace(uname, scope), 'op': 'Parameter', 'inputs': [], 'attr': str(n.type())}) + + return nodes + + +def graph(model, args, verbose=False): + import torch + with torch.onnx.set_training(model, False): + trace, _ = torch.jit.trace(model, args) + if LooseVersion(torch.__version__) >= LooseVersion("0.4"): + torch.onnx._optimize_trace(trace, False) + else: + torch.onnx._optimize_trace(trace) + graph = trace.graph() + if verbose: + print(graph) + list_of_nodes = parse(graph) + nodes = [] + for node in list_of_nodes: + nodes.append( + NodeDef(name=node['name'], op=node['op'], input=node['inputs'], + attr={'lanpa': AttrValue(s=node['attr'].encode(encoding='utf_8'))})) + return GraphDef(node=nodes, versions=VersionDef(producer=22)) diff --git a/tensorboardX/graph_onnx.py b/tensorboardX/graph_onnx.py new file mode 100644 index 0000000..96a3534 --- /dev/null +++ b/tensorboardX/graph_onnx.py @@ -0,0 +1,158 @@ +from .src.graph_pb2 import GraphDef +from .src.node_def_pb2 import NodeDef +from .src.versions_pb2 import VersionDef +from .src.attr_value_pb2 import AttrValue +from .src.tensor_shape_pb2 import TensorShapeProto +# from .src.onnx_pb2 import ModelProto + + +def gg(fname): + import onnx # 0.2.1 + m = onnx.load(fname) + nodes_proto = [] + nodes = [] + g = m.graph + import itertools + for node in itertools.chain(g.input, g.output): + nodes_proto.append(node) + + for node in nodes_proto: + shapeproto = TensorShapeProto( + dim=[TensorShapeProto.Dim(size=d.dim_value) for d in node.type.tensor_type.shape.dim]) + nodes.append(NodeDef( + name=node.name, + op='Variable', + input=[], + attr={ + 'dtype': AttrValue(type=node.type.tensor_type.elem_type), + 'shape': AttrValue(shape=shapeproto), + }) + ) + + for node in g.node: + attr = [] + for s in node.attribute: + attr.append(' = '.join([str(f[1]) for f in s.ListFields()])) + attr = ', '.join(attr).encode(encoding='utf_8') + + nodes.append(NodeDef( + name=node.output[0], + op=node.op_type, + input=node.input, + attr={'parameters': AttrValue(s=attr)}, + )) + # two pass token replacement, appends opname to object id + mapping = {} + for node in nodes: + mapping[node.name] = node.op + '_' + node.name + + nodes, mapping = updatenodes(nodes, mapping) + mapping = smartGrouping(nodes, mapping) + nodes, mapping = updatenodes(nodes, mapping) + + return GraphDef(node=nodes, versions=VersionDef(producer=22)) + + +def updatenodes(nodes, mapping): + for node in nodes: + newname = mapping[node.name] + node.name = newname + newinput = [] + for inputnode in list(node.input): + newinput.append(mapping[inputnode]) + node.input.remove(inputnode) + node.input.extend(newinput) + newmap = {} + for k, v in mapping.items(): + newmap[v] = v + return nodes, newmap + + +def findnode(nodes, name): + """ input: node name + returns: node object + """ + for n in nodes: + if n.name == name: + return n + + +def parser(s, nodes, node): + print(s) + if len(s) == 0: + return + if len(s) > 0: + if s[0] == node.op: + print(s[0], node.name, s[1], node.input) + for n in node.input: + print(n, s[1]) + parser(s[1], nodes, findnode(nodes, n)) + else: + return False + + +# TODO: use recursive parse + +def smartGrouping(nodes, mapping): + # a Fully Conv is: (TODO: check var1.size(0)==var2.size(0)) + # GEMM <-- Variable (c1) + # ^-- Transpose (c2) <-- Variable (c3) + + # a Conv with bias is: (TODO: check var1.size(0)==var2.size(0)) + # Add <-- Conv (c2) <-- Variable (c3) + # ^-- Variable (c1) + # + # gemm = ('Gemm', ('Variable', ('Transpose', ('Variable')))) + + FCcounter = 1 + Convcounter = 1 + for node in nodes: + if node.op == 'Gemm': + c1 = c2 = c3 = False + for name_in in node.input: + n = findnode(nodes, name_in) + if n.op == 'Variable': + c1 = True + c1name = n.name + if n.op == 'Transpose': + c2 = True + c2name = n.name + if len(n.input) == 1: + nn = findnode(nodes, n.input[0]) + if nn.op == 'Variable': + c3 = True + c3name = nn.name + # print(n.op, n.name, c1, c2, c3) + if c1 and c2 and c3: + # print(c1name, c2name, c3name) + mapping[c1name] = 'FC{}/{}'.format(FCcounter, c1name) + mapping[c2name] = 'FC{}/{}'.format(FCcounter, c2name) + mapping[c3name] = 'FC{}/{}'.format(FCcounter, c3name) + mapping[node.name] = 'FC{}/{}'.format(FCcounter, node.name) + FCcounter += 1 + continue + if node.op == 'Add': + c1 = c2 = c3 = False + for name_in in node.input: + n = findnode(nodes, name_in) + if n.op == 'Variable': + c1 = True + c1name = n.name + if n.op == 'Conv': + c2 = True + c2name = n.name + if len(n.input) >= 1: + for nn_name in n.input: + nn = findnode(nodes, nn_name) + if nn.op == 'Variable': + c3 = True + c3name = nn.name + + if c1 and c2 and c3: + # print(c1name, c2name, c3name) + mapping[c1name] = 'Conv{}/{}'.format(Convcounter, c1name) + mapping[c2name] = 'Conv{}/{}'.format(Convcounter, c2name) + mapping[c3name] = 'Conv{}/{}'.format(Convcounter, c3name) + mapping[node.name] = 'Conv{}/{}'.format(Convcounter, node.name) + Convcounter += 1 + return mapping diff --git a/tensorboardX/record_writer.py b/tensorboardX/record_writer.py new file mode 100644 index 0000000..c14bd7b --- /dev/null +++ b/tensorboardX/record_writer.py @@ -0,0 +1,47 @@ +""" +To write tf_record into file. Here we use it for tensorboard's event writting. +The code was borrow from https://github.com/TeamHG-Memex/tensorboard_logger +""" + +import re +import struct + +from .crc32c import crc32c + +_VALID_OP_NAME_START = re.compile('^[A-Za-z0-9.]') +_VALID_OP_NAME_PART = re.compile('[A-Za-z0-9_.\\-/]+') + + +class RecordWriter(object): + def __init__(self, path, flush_secs=2): + self._name_to_tf_name = {} + self._tf_names = set() + self.path = path + self.flush_secs = flush_secs # TODO. flush every flush_secs, not every time. + self._writer = None + self._writer = open(path, 'wb') + + def write(self, event_str): + w = self._writer.write + header = struct.pack('Q', len(event_str)) + w(header) + w(struct.pack('I', masked_crc32c(header))) + w(event_str) + w(struct.pack('I', masked_crc32c(event_str))) + self._writer.flush() + + +def masked_crc32c(data): + x = u32(crc32c(data)) + return u32(((x >> 15) | u32(x << 17)) + 0xa282ead8) + + +def u32(x): + return x & 0xffffffff + + +def make_valid_tf_name(name): + if not _VALID_OP_NAME_START.match(name): + # Must make it valid somehow, but don't want to remove stuff + name = '.' + name + return '_'.join(_VALID_OP_NAME_PART.findall(name)) diff --git a/tensorboardX/src/__init__.py b/tensorboardX/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tensorboardX/src/attr_value_pb2.py b/tensorboardX/src/attr_value_pb2.py new file mode 100644 index 0000000..95fae04 --- /dev/null +++ b/tensorboardX/src/attr_value_pb2.py @@ -0,0 +1,372 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/attr_value.proto + +import sys +_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +from google.protobuf import descriptor_pb2 +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from tensorboardX.src import tensor_pb2 as tensorboardX_dot_src_dot_tensor__pb2 +from tensorboardX.src import tensor_shape_pb2 as tensorboardX_dot_src_dot_tensor__shape__pb2 +from tensorboardX.src import types_pb2 as tensorboardX_dot_src_dot_types__pb2 + + +DESCRIPTOR = _descriptor.FileDescriptor( + name='tensorboardX/src/attr_value.proto', + package='tensorboard', + syntax='proto3', + serialized_pb=_b('\n!tensorboardX/src/attr_value.proto\x12\x0btensorboard\x1a\x1dtensorboardX/src/tensor.proto\x1a#tensorboardX/src/tensor_shape.proto\x1a\x1ctensorboardX/src/types.proto\"\xaf\x04\n\tAttrValue\x12\x0b\n\x01s\x18\x02 \x01(\x0cH\x00\x12\x0b\n\x01i\x18\x03 \x01(\x03H\x00\x12\x0b\n\x01\x66\x18\x04 \x01(\x02H\x00\x12\x0b\n\x01\x62\x18\x05 \x01(\x08H\x00\x12%\n\x04type\x18\x06 \x01(\x0e\x32\x15.tensorboard.DataTypeH\x00\x12.\n\x05shape\x18\x07 \x01(\x0b\x32\x1d.tensorboard.TensorShapeProtoH\x00\x12*\n\x06tensor\x18\x08 \x01(\x0b\x32\x18.tensorboard.TensorProtoH\x00\x12\x30\n\x04list\x18\x01 \x01(\x0b\x32 .tensorboard.AttrValue.ListValueH\x00\x12)\n\x04\x66unc\x18\n \x01(\x0b\x32\x19.tensorboard.NameAttrListH\x00\x12\x15\n\x0bplaceholder\x18\t \x01(\tH\x00\x1a\xed\x01\n\tListValue\x12\t\n\x01s\x18\x02 \x03(\x0c\x12\r\n\x01i\x18\x03 \x03(\x03\x42\x02\x10\x01\x12\r\n\x01\x66\x18\x04 \x03(\x02\x42\x02\x10\x01\x12\r\n\x01\x62\x18\x05 \x03(\x08\x42\x02\x10\x01\x12\'\n\x04type\x18\x06 \x03(\x0e\x32\x15.tensorboard.DataTypeB\x02\x10\x01\x12,\n\x05shape\x18\x07 \x03(\x0b\x32\x1d.tensorboard.TensorShapeProto\x12(\n\x06tensor\x18\x08 \x03(\x0b\x32\x18.tensorboard.TensorProto\x12\'\n\x04\x66unc\x18\t \x03(\x0b\x32\x19.tensorboard.NameAttrListB\x07\n\x05value\"\x94\x01\n\x0cNameAttrList\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x31\n\x04\x61ttr\x18\x02 \x03(\x0b\x32#.tensorboard.NameAttrList.AttrEntry\x1a\x43\n\tAttrEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12%\n\x05value\x18\x02 \x01(\x0b\x32\x16.tensorboard.AttrValue:\x02\x38\x01\x42\x30\n\x18org.tensorflow.frameworkB\x0f\x41ttrValueProtosP\x01\xf8\x01\x01\x62\x06proto3') + , + dependencies=[tensorboardX_dot_src_dot_tensor__pb2.DESCRIPTOR,tensorboardX_dot_src_dot_tensor__shape__pb2.DESCRIPTOR,tensorboardX_dot_src_dot_types__pb2.DESCRIPTOR,]) + + + + +_ATTRVALUE_LISTVALUE = _descriptor.Descriptor( + name='ListValue', + full_name='tensorboard.AttrValue.ListValue', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='s', full_name='tensorboard.AttrValue.ListValue.s', index=0, + number=2, type=12, cpp_type=9, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='i', full_name='tensorboard.AttrValue.ListValue.i', index=1, + number=3, type=3, cpp_type=2, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))), + _descriptor.FieldDescriptor( + name='f', full_name='tensorboard.AttrValue.ListValue.f', index=2, + number=4, type=2, cpp_type=6, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))), + _descriptor.FieldDescriptor( + name='b', full_name='tensorboard.AttrValue.ListValue.b', index=3, + number=5, type=8, cpp_type=7, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))), + _descriptor.FieldDescriptor( + name='type', full_name='tensorboard.AttrValue.ListValue.type', index=4, + number=6, type=14, cpp_type=8, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))), + _descriptor.FieldDescriptor( + name='shape', full_name='tensorboard.AttrValue.ListValue.shape', index=5, + number=7, type=11, cpp_type=10, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='tensor', full_name='tensorboard.AttrValue.ListValue.tensor', index=6, + number=8, type=11, cpp_type=10, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='func', full_name='tensorboard.AttrValue.ListValue.func', index=7, + number=9, type=11, cpp_type=10, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=462, + serialized_end=699, +) + +_ATTRVALUE = _descriptor.Descriptor( + name='AttrValue', + full_name='tensorboard.AttrValue', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='s', full_name='tensorboard.AttrValue.s', index=0, + number=2, type=12, cpp_type=9, label=1, + has_default_value=False, default_value=_b(""), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='i', full_name='tensorboard.AttrValue.i', index=1, + number=3, type=3, cpp_type=2, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='f', full_name='tensorboard.AttrValue.f', index=2, + number=4, type=2, cpp_type=6, label=1, + has_default_value=False, default_value=float(0), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='b', full_name='tensorboard.AttrValue.b', index=3, + number=5, type=8, cpp_type=7, label=1, + has_default_value=False, default_value=False, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='type', full_name='tensorboard.AttrValue.type', index=4, + number=6, type=14, cpp_type=8, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='shape', full_name='tensorboard.AttrValue.shape', index=5, + number=7, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='tensor', full_name='tensorboard.AttrValue.tensor', index=6, + number=8, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='list', full_name='tensorboard.AttrValue.list', index=7, + number=1, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='func', full_name='tensorboard.AttrValue.func', index=8, + number=10, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='placeholder', full_name='tensorboard.AttrValue.placeholder', index=9, + number=9, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[_ATTRVALUE_LISTVALUE, ], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + _descriptor.OneofDescriptor( + name='value', full_name='tensorboard.AttrValue.value', + index=0, containing_type=None, fields=[]), + ], + serialized_start=149, + serialized_end=708, +) + + +_NAMEATTRLIST_ATTRENTRY = _descriptor.Descriptor( + name='AttrEntry', + full_name='tensorboard.NameAttrList.AttrEntry', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='key', full_name='tensorboard.NameAttrList.AttrEntry.key', index=0, + number=1, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='value', full_name='tensorboard.NameAttrList.AttrEntry.value', index=1, + number=2, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=_descriptor._ParseOptions(descriptor_pb2.MessageOptions(), _b('8\001')), + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=792, + serialized_end=859, +) + +_NAMEATTRLIST = _descriptor.Descriptor( + name='NameAttrList', + full_name='tensorboard.NameAttrList', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='name', full_name='tensorboard.NameAttrList.name', index=0, + number=1, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='attr', full_name='tensorboard.NameAttrList.attr', index=1, + number=2, type=11, cpp_type=10, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[_NAMEATTRLIST_ATTRENTRY, ], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=711, + serialized_end=859, +) + +_ATTRVALUE_LISTVALUE.fields_by_name['type'].enum_type = tensorboardX_dot_src_dot_types__pb2._DATATYPE +_ATTRVALUE_LISTVALUE.fields_by_name['shape'].message_type = tensorboardX_dot_src_dot_tensor__shape__pb2._TENSORSHAPEPROTO +_ATTRVALUE_LISTVALUE.fields_by_name['tensor'].message_type = tensorboardX_dot_src_dot_tensor__pb2._TENSORPROTO +_ATTRVALUE_LISTVALUE.fields_by_name['func'].message_type = _NAMEATTRLIST +_ATTRVALUE_LISTVALUE.containing_type = _ATTRVALUE +_ATTRVALUE.fields_by_name['type'].enum_type = tensorboardX_dot_src_dot_types__pb2._DATATYPE +_ATTRVALUE.fields_by_name['shape'].message_type = tensorboardX_dot_src_dot_tensor__shape__pb2._TENSORSHAPEPROTO +_ATTRVALUE.fields_by_name['tensor'].message_type = tensorboardX_dot_src_dot_tensor__pb2._TENSORPROTO +_ATTRVALUE.fields_by_name['list'].message_type = _ATTRVALUE_LISTVALUE +_ATTRVALUE.fields_by_name['func'].message_type = _NAMEATTRLIST +_ATTRVALUE.oneofs_by_name['value'].fields.append( + _ATTRVALUE.fields_by_name['s']) +_ATTRVALUE.fields_by_name['s'].containing_oneof = _ATTRVALUE.oneofs_by_name['value'] +_ATTRVALUE.oneofs_by_name['value'].fields.append( + _ATTRVALUE.fields_by_name['i']) +_ATTRVALUE.fields_by_name['i'].containing_oneof = _ATTRVALUE.oneofs_by_name['value'] +_ATTRVALUE.oneofs_by_name['value'].fields.append( + _ATTRVALUE.fields_by_name['f']) +_ATTRVALUE.fields_by_name['f'].containing_oneof = _ATTRVALUE.oneofs_by_name['value'] +_ATTRVALUE.oneofs_by_name['value'].fields.append( + _ATTRVALUE.fields_by_name['b']) +_ATTRVALUE.fields_by_name['b'].containing_oneof = _ATTRVALUE.oneofs_by_name['value'] +_ATTRVALUE.oneofs_by_name['value'].fields.append( + _ATTRVALUE.fields_by_name['type']) +_ATTRVALUE.fields_by_name['type'].containing_oneof = _ATTRVALUE.oneofs_by_name['value'] +_ATTRVALUE.oneofs_by_name['value'].fields.append( + _ATTRVALUE.fields_by_name['shape']) +_ATTRVALUE.fields_by_name['shape'].containing_oneof = _ATTRVALUE.oneofs_by_name['value'] +_ATTRVALUE.oneofs_by_name['value'].fields.append( + _ATTRVALUE.fields_by_name['tensor']) +_ATTRVALUE.fields_by_name['tensor'].containing_oneof = _ATTRVALUE.oneofs_by_name['value'] +_ATTRVALUE.oneofs_by_name['value'].fields.append( + _ATTRVALUE.fields_by_name['list']) +_ATTRVALUE.fields_by_name['list'].containing_oneof = _ATTRVALUE.oneofs_by_name['value'] +_ATTRVALUE.oneofs_by_name['value'].fields.append( + _ATTRVALUE.fields_by_name['func']) +_ATTRVALUE.fields_by_name['func'].containing_oneof = _ATTRVALUE.oneofs_by_name['value'] +_ATTRVALUE.oneofs_by_name['value'].fields.append( + _ATTRVALUE.fields_by_name['placeholder']) +_ATTRVALUE.fields_by_name['placeholder'].containing_oneof = _ATTRVALUE.oneofs_by_name['value'] +_NAMEATTRLIST_ATTRENTRY.fields_by_name['value'].message_type = _ATTRVALUE +_NAMEATTRLIST_ATTRENTRY.containing_type = _NAMEATTRLIST +_NAMEATTRLIST.fields_by_name['attr'].message_type = _NAMEATTRLIST_ATTRENTRY +DESCRIPTOR.message_types_by_name['AttrValue'] = _ATTRVALUE +DESCRIPTOR.message_types_by_name['NameAttrList'] = _NAMEATTRLIST +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +AttrValue = _reflection.GeneratedProtocolMessageType('AttrValue', (_message.Message,), dict( + + ListValue = _reflection.GeneratedProtocolMessageType('ListValue', (_message.Message,), dict( + DESCRIPTOR = _ATTRVALUE_LISTVALUE, + __module__ = 'tensorboardX.src.attr_value_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.AttrValue.ListValue) + )) + , + DESCRIPTOR = _ATTRVALUE, + __module__ = 'tensorboardX.src.attr_value_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.AttrValue) + )) +_sym_db.RegisterMessage(AttrValue) +_sym_db.RegisterMessage(AttrValue.ListValue) + +NameAttrList = _reflection.GeneratedProtocolMessageType('NameAttrList', (_message.Message,), dict( + + AttrEntry = _reflection.GeneratedProtocolMessageType('AttrEntry', (_message.Message,), dict( + DESCRIPTOR = _NAMEATTRLIST_ATTRENTRY, + __module__ = 'tensorboardX.src.attr_value_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.NameAttrList.AttrEntry) + )) + , + DESCRIPTOR = _NAMEATTRLIST, + __module__ = 'tensorboardX.src.attr_value_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.NameAttrList) + )) +_sym_db.RegisterMessage(NameAttrList) +_sym_db.RegisterMessage(NameAttrList.AttrEntry) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\017AttrValueProtosP\001\370\001\001')) +_ATTRVALUE_LISTVALUE.fields_by_name['i'].has_options = True +_ATTRVALUE_LISTVALUE.fields_by_name['i']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001')) +_ATTRVALUE_LISTVALUE.fields_by_name['f'].has_options = True +_ATTRVALUE_LISTVALUE.fields_by_name['f']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001')) +_ATTRVALUE_LISTVALUE.fields_by_name['b'].has_options = True +_ATTRVALUE_LISTVALUE.fields_by_name['b']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001')) +_ATTRVALUE_LISTVALUE.fields_by_name['type'].has_options = True +_ATTRVALUE_LISTVALUE.fields_by_name['type']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001')) +_NAMEATTRLIST_ATTRENTRY.has_options = True +_NAMEATTRLIST_ATTRENTRY._options = _descriptor._ParseOptions(descriptor_pb2.MessageOptions(), _b('8\001')) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/event_pb2.py b/tensorboardX/src/event_pb2.py new file mode 100644 index 0000000..459dfbd --- /dev/null +++ b/tensorboardX/src/event_pb2.py @@ -0,0 +1,376 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/event.proto + +import sys +_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +from google.protobuf import descriptor_pb2 +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from tensorboardX.src import summary_pb2 as tensorboardX_dot_src_dot_summary__pb2 + + +DESCRIPTOR = _descriptor.FileDescriptor( + name='tensorboardX/src/event.proto', + package='tensorboard', + syntax='proto3', + serialized_pb=_b('\n\x1ctensorboardX/src/event.proto\x12\x0btensorboard\x1a\x1etensorboardX/src/summary.proto\"\xbf\x02\n\x05\x45vent\x12\x11\n\twall_time\x18\x01 \x01(\x01\x12\x0c\n\x04step\x18\x02 \x01(\x03\x12\x16\n\x0c\x66ile_version\x18\x03 \x01(\tH\x00\x12\x13\n\tgraph_def\x18\x04 \x01(\x0cH\x00\x12\'\n\x07summary\x18\x05 \x01(\x0b\x32\x14.tensorboard.SummaryH\x00\x12.\n\x0blog_message\x18\x06 \x01(\x0b\x32\x17.tensorboard.LogMessageH\x00\x12.\n\x0bsession_log\x18\x07 \x01(\x0b\x32\x17.tensorboard.SessionLogH\x00\x12=\n\x13tagged_run_metadata\x18\x08 \x01(\x0b\x32\x1e.tensorboard.TaggedRunMetadataH\x00\x12\x18\n\x0emeta_graph_def\x18\t \x01(\x0cH\x00\x42\x06\n\x04what\"\x96\x01\n\nLogMessage\x12,\n\x05level\x18\x01 \x01(\x0e\x32\x1d.tensorboard.LogMessage.Level\x12\x0f\n\x07message\x18\x02 \x01(\t\"I\n\x05Level\x12\x0b\n\x07UNKNOWN\x10\x00\x12\t\n\x05\x44\x45\x42UG\x10\n\x12\x08\n\x04INFO\x10\x14\x12\x08\n\x04WARN\x10\x1e\x12\t\n\x05\x45RROR\x10(\x12\t\n\x05\x46\x41TAL\x10\x32\"\xb7\x01\n\nSessionLog\x12\x35\n\x06status\x18\x01 \x01(\x0e\x32%.tensorboard.SessionLog.SessionStatus\x12\x17\n\x0f\x63heckpoint_path\x18\x02 \x01(\t\x12\x0b\n\x03msg\x18\x03 \x01(\t\"L\n\rSessionStatus\x12\x16\n\x12STATUS_UNSPECIFIED\x10\x00\x12\t\n\x05START\x10\x01\x12\x08\n\x04STOP\x10\x02\x12\x0e\n\nCHECKPOINT\x10\x03\"6\n\x11TaggedRunMetadata\x12\x0b\n\x03tag\x18\x01 \x01(\t\x12\x14\n\x0crun_metadata\x18\x02 \x01(\x0c\x42\'\n\x13org.tensorflow.utilB\x0b\x45ventProtosP\x01\xf8\x01\x01\x62\x06proto3') + , + dependencies=[tensorboardX_dot_src_dot_summary__pb2.DESCRIPTOR,]) + + + +_LOGMESSAGE_LEVEL = _descriptor.EnumDescriptor( + name='Level', + full_name='tensorboard.LogMessage.Level', + filename=None, + file=DESCRIPTOR, + values=[ + _descriptor.EnumValueDescriptor( + name='UNKNOWN', index=0, number=0, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DEBUG', index=1, number=10, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='INFO', index=2, number=20, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='WARN', index=3, number=30, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='ERROR', index=4, number=40, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='FATAL', index=5, number=50, + options=None, + type=None), + ], + containing_type=None, + options=None, + serialized_start=477, + serialized_end=550, +) +_sym_db.RegisterEnumDescriptor(_LOGMESSAGE_LEVEL) + +_SESSIONLOG_SESSIONSTATUS = _descriptor.EnumDescriptor( + name='SessionStatus', + full_name='tensorboard.SessionLog.SessionStatus', + filename=None, + file=DESCRIPTOR, + values=[ + _descriptor.EnumValueDescriptor( + name='STATUS_UNSPECIFIED', index=0, number=0, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='START', index=1, number=1, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='STOP', index=2, number=2, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='CHECKPOINT', index=3, number=3, + options=None, + type=None), + ], + containing_type=None, + options=None, + serialized_start=660, + serialized_end=736, +) +_sym_db.RegisterEnumDescriptor(_SESSIONLOG_SESSIONSTATUS) + + +_EVENT = _descriptor.Descriptor( + name='Event', + full_name='tensorboard.Event', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='wall_time', full_name='tensorboard.Event.wall_time', index=0, + number=1, type=1, cpp_type=5, label=1, + has_default_value=False, default_value=float(0), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='step', full_name='tensorboard.Event.step', index=1, + number=2, type=3, cpp_type=2, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='file_version', full_name='tensorboard.Event.file_version', index=2, + number=3, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='graph_def', full_name='tensorboard.Event.graph_def', index=3, + number=4, type=12, cpp_type=9, label=1, + has_default_value=False, default_value=_b(""), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='summary', full_name='tensorboard.Event.summary', index=4, + number=5, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='log_message', full_name='tensorboard.Event.log_message', index=5, + number=6, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='session_log', full_name='tensorboard.Event.session_log', index=6, + number=7, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='tagged_run_metadata', full_name='tensorboard.Event.tagged_run_metadata', index=7, + number=8, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='meta_graph_def', full_name='tensorboard.Event.meta_graph_def', index=8, + number=9, type=12, cpp_type=9, label=1, + has_default_value=False, default_value=_b(""), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + _descriptor.OneofDescriptor( + name='what', full_name='tensorboard.Event.what', + index=0, containing_type=None, fields=[]), + ], + serialized_start=78, + serialized_end=397, +) + + +_LOGMESSAGE = _descriptor.Descriptor( + name='LogMessage', + full_name='tensorboard.LogMessage', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='level', full_name='tensorboard.LogMessage.level', index=0, + number=1, type=14, cpp_type=8, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='message', full_name='tensorboard.LogMessage.message', index=1, + number=2, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + _LOGMESSAGE_LEVEL, + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=400, + serialized_end=550, +) + + +_SESSIONLOG = _descriptor.Descriptor( + name='SessionLog', + full_name='tensorboard.SessionLog', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='status', full_name='tensorboard.SessionLog.status', index=0, + number=1, type=14, cpp_type=8, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='checkpoint_path', full_name='tensorboard.SessionLog.checkpoint_path', index=1, + number=2, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='msg', full_name='tensorboard.SessionLog.msg', index=2, + number=3, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + _SESSIONLOG_SESSIONSTATUS, + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=553, + serialized_end=736, +) + + +_TAGGEDRUNMETADATA = _descriptor.Descriptor( + name='TaggedRunMetadata', + full_name='tensorboard.TaggedRunMetadata', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='tag', full_name='tensorboard.TaggedRunMetadata.tag', index=0, + number=1, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='run_metadata', full_name='tensorboard.TaggedRunMetadata.run_metadata', index=1, + number=2, type=12, cpp_type=9, label=1, + has_default_value=False, default_value=_b(""), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=738, + serialized_end=792, +) + +_EVENT.fields_by_name['summary'].message_type = tensorboardX_dot_src_dot_summary__pb2._SUMMARY +_EVENT.fields_by_name['log_message'].message_type = _LOGMESSAGE +_EVENT.fields_by_name['session_log'].message_type = _SESSIONLOG +_EVENT.fields_by_name['tagged_run_metadata'].message_type = _TAGGEDRUNMETADATA +_EVENT.oneofs_by_name['what'].fields.append( + _EVENT.fields_by_name['file_version']) +_EVENT.fields_by_name['file_version'].containing_oneof = _EVENT.oneofs_by_name['what'] +_EVENT.oneofs_by_name['what'].fields.append( + _EVENT.fields_by_name['graph_def']) +_EVENT.fields_by_name['graph_def'].containing_oneof = _EVENT.oneofs_by_name['what'] +_EVENT.oneofs_by_name['what'].fields.append( + _EVENT.fields_by_name['summary']) +_EVENT.fields_by_name['summary'].containing_oneof = _EVENT.oneofs_by_name['what'] +_EVENT.oneofs_by_name['what'].fields.append( + _EVENT.fields_by_name['log_message']) +_EVENT.fields_by_name['log_message'].containing_oneof = _EVENT.oneofs_by_name['what'] +_EVENT.oneofs_by_name['what'].fields.append( + _EVENT.fields_by_name['session_log']) +_EVENT.fields_by_name['session_log'].containing_oneof = _EVENT.oneofs_by_name['what'] +_EVENT.oneofs_by_name['what'].fields.append( + _EVENT.fields_by_name['tagged_run_metadata']) +_EVENT.fields_by_name['tagged_run_metadata'].containing_oneof = _EVENT.oneofs_by_name['what'] +_EVENT.oneofs_by_name['what'].fields.append( + _EVENT.fields_by_name['meta_graph_def']) +_EVENT.fields_by_name['meta_graph_def'].containing_oneof = _EVENT.oneofs_by_name['what'] +_LOGMESSAGE.fields_by_name['level'].enum_type = _LOGMESSAGE_LEVEL +_LOGMESSAGE_LEVEL.containing_type = _LOGMESSAGE +_SESSIONLOG.fields_by_name['status'].enum_type = _SESSIONLOG_SESSIONSTATUS +_SESSIONLOG_SESSIONSTATUS.containing_type = _SESSIONLOG +DESCRIPTOR.message_types_by_name['Event'] = _EVENT +DESCRIPTOR.message_types_by_name['LogMessage'] = _LOGMESSAGE +DESCRIPTOR.message_types_by_name['SessionLog'] = _SESSIONLOG +DESCRIPTOR.message_types_by_name['TaggedRunMetadata'] = _TAGGEDRUNMETADATA +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +Event = _reflection.GeneratedProtocolMessageType('Event', (_message.Message,), dict( + DESCRIPTOR = _EVENT, + __module__ = 'tensorboardX.src.event_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.Event) + )) +_sym_db.RegisterMessage(Event) + +LogMessage = _reflection.GeneratedProtocolMessageType('LogMessage', (_message.Message,), dict( + DESCRIPTOR = _LOGMESSAGE, + __module__ = 'tensorboardX.src.event_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.LogMessage) + )) +_sym_db.RegisterMessage(LogMessage) + +SessionLog = _reflection.GeneratedProtocolMessageType('SessionLog', (_message.Message,), dict( + DESCRIPTOR = _SESSIONLOG, + __module__ = 'tensorboardX.src.event_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.SessionLog) + )) +_sym_db.RegisterMessage(SessionLog) + +TaggedRunMetadata = _reflection.GeneratedProtocolMessageType('TaggedRunMetadata', (_message.Message,), dict( + DESCRIPTOR = _TAGGEDRUNMETADATA, + __module__ = 'tensorboardX.src.event_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.TaggedRunMetadata) + )) +_sym_db.RegisterMessage(TaggedRunMetadata) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\023org.tensorflow.utilB\013EventProtosP\001\370\001\001')) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/graph_pb2.py b/tensorboardX/src/graph_pb2.py new file mode 100644 index 0000000..23dda00 --- /dev/null +++ b/tensorboardX/src/graph_pb2.py @@ -0,0 +1,92 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/graph.proto + +import sys +_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +from google.protobuf import descriptor_pb2 +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from tensorboardX.src import node_def_pb2 as tensorboardX_dot_src_dot_node__def__pb2 +from tensorboardX.src import versions_pb2 as tensorboardX_dot_src_dot_versions__pb2 + + +DESCRIPTOR = _descriptor.FileDescriptor( + name='tensorboardX/src/graph.proto', + package='tensorboard', + syntax='proto3', + serialized_pb=_b('\n\x1ctensorboardX/src/graph.proto\x12\x0btensorboard\x1a\x1ftensorboardX/src/node_def.proto\x1a\x1ftensorboardX/src/versions.proto\"n\n\x08GraphDef\x12\"\n\x04node\x18\x01 \x03(\x0b\x32\x14.tensorboard.NodeDef\x12)\n\x08versions\x18\x04 \x01(\x0b\x32\x17.tensorboard.VersionDef\x12\x13\n\x07version\x18\x03 \x01(\x05\x42\x02\x18\x01\x42,\n\x18org.tensorflow.frameworkB\x0bGraphProtosP\x01\xf8\x01\x01\x62\x06proto3') + , + dependencies=[tensorboardX_dot_src_dot_node__def__pb2.DESCRIPTOR,tensorboardX_dot_src_dot_versions__pb2.DESCRIPTOR,]) + + + + +_GRAPHDEF = _descriptor.Descriptor( + name='GraphDef', + full_name='tensorboard.GraphDef', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='node', full_name='tensorboard.GraphDef.node', index=0, + number=1, type=11, cpp_type=10, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='versions', full_name='tensorboard.GraphDef.versions', index=1, + number=4, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='version', full_name='tensorboard.GraphDef.version', index=2, + number=3, type=5, cpp_type=1, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\030\001'))), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=111, + serialized_end=221, +) + +_GRAPHDEF.fields_by_name['node'].message_type = tensorboardX_dot_src_dot_node__def__pb2._NODEDEF +_GRAPHDEF.fields_by_name['versions'].message_type = tensorboardX_dot_src_dot_versions__pb2._VERSIONDEF +DESCRIPTOR.message_types_by_name['GraphDef'] = _GRAPHDEF +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +GraphDef = _reflection.GeneratedProtocolMessageType('GraphDef', (_message.Message,), dict( + DESCRIPTOR = _GRAPHDEF, + __module__ = 'tensorboardX.src.graph_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.GraphDef) + )) +_sym_db.RegisterMessage(GraphDef) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\013GraphProtosP\001\370\001\001')) +_GRAPHDEF.fields_by_name['version'].has_options = True +_GRAPHDEF.fields_by_name['version']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\030\001')) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/node_def_pb2.py b/tensorboardX/src/node_def_pb2.py new file mode 100644 index 0000000..11ec64e --- /dev/null +++ b/tensorboardX/src/node_def_pb2.py @@ -0,0 +1,151 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/node_def.proto + +import sys +_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +from google.protobuf import descriptor_pb2 +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from tensorboardX.src import attr_value_pb2 as tensorboardX_dot_src_dot_attr__value__pb2 + + +DESCRIPTOR = _descriptor.FileDescriptor( + name='tensorboardX/src/node_def.proto', + package='tensorboard', + syntax='proto3', + serialized_pb=_b('\n\x1ftensorboardX/src/node_def.proto\x12\x0btensorboard\x1a!tensorboardX/src/attr_value.proto\"\xb5\x01\n\x07NodeDef\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\n\n\x02op\x18\x02 \x01(\t\x12\r\n\x05input\x18\x03 \x03(\t\x12\x0e\n\x06\x64\x65vice\x18\x04 \x01(\t\x12,\n\x04\x61ttr\x18\x05 \x03(\x0b\x32\x1e.tensorboard.NodeDef.AttrEntry\x1a\x43\n\tAttrEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12%\n\x05value\x18\x02 \x01(\x0b\x32\x16.tensorboard.AttrValue:\x02\x38\x01\x42*\n\x18org.tensorflow.frameworkB\tNodeProtoP\x01\xf8\x01\x01\x62\x06proto3') + , + dependencies=[tensorboardX_dot_src_dot_attr__value__pb2.DESCRIPTOR,]) + + + + +_NODEDEF_ATTRENTRY = _descriptor.Descriptor( + name='AttrEntry', + full_name='tensorboard.NodeDef.AttrEntry', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='key', full_name='tensorboard.NodeDef.AttrEntry.key', index=0, + number=1, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='value', full_name='tensorboard.NodeDef.AttrEntry.value', index=1, + number=2, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=_descriptor._ParseOptions(descriptor_pb2.MessageOptions(), _b('8\001')), + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=198, + serialized_end=265, +) + +_NODEDEF = _descriptor.Descriptor( + name='NodeDef', + full_name='tensorboard.NodeDef', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='name', full_name='tensorboard.NodeDef.name', index=0, + number=1, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='op', full_name='tensorboard.NodeDef.op', index=1, + number=2, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='input', full_name='tensorboard.NodeDef.input', index=2, + number=3, type=9, cpp_type=9, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='device', full_name='tensorboard.NodeDef.device', index=3, + number=4, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='attr', full_name='tensorboard.NodeDef.attr', index=4, + number=5, type=11, cpp_type=10, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[_NODEDEF_ATTRENTRY, ], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=84, + serialized_end=265, +) + +_NODEDEF_ATTRENTRY.fields_by_name['value'].message_type = tensorboardX_dot_src_dot_attr__value__pb2._ATTRVALUE +_NODEDEF_ATTRENTRY.containing_type = _NODEDEF +_NODEDEF.fields_by_name['attr'].message_type = _NODEDEF_ATTRENTRY +DESCRIPTOR.message_types_by_name['NodeDef'] = _NODEDEF +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +NodeDef = _reflection.GeneratedProtocolMessageType('NodeDef', (_message.Message,), dict( + + AttrEntry = _reflection.GeneratedProtocolMessageType('AttrEntry', (_message.Message,), dict( + DESCRIPTOR = _NODEDEF_ATTRENTRY, + __module__ = 'tensorboardX.src.node_def_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.NodeDef.AttrEntry) + )) + , + DESCRIPTOR = _NODEDEF, + __module__ = 'tensorboardX.src.node_def_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.NodeDef) + )) +_sym_db.RegisterMessage(NodeDef) +_sym_db.RegisterMessage(NodeDef.AttrEntry) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\tNodeProtoP\001\370\001\001')) +_NODEDEF_ATTRENTRY.has_options = True +_NODEDEF_ATTRENTRY._options = _descriptor._ParseOptions(descriptor_pb2.MessageOptions(), _b('8\001')) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/plugin_pr_curve_pb2.py b/tensorboardX/src/plugin_pr_curve_pb2.py new file mode 100644 index 0000000..c219d52 --- /dev/null +++ b/tensorboardX/src/plugin_pr_curve_pb2.py @@ -0,0 +1,76 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/plugin_pr_curve.proto + +import sys +_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +from google.protobuf import descriptor_pb2 +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor.FileDescriptor( + name='tensorboardX/src/plugin_pr_curve.proto', + package='tensorboard', + syntax='proto3', + serialized_pb=_b('\n&tensorboardX/src/plugin_pr_curve.proto\x12\x0btensorboard\"<\n\x11PrCurvePluginData\x12\x0f\n\x07version\x18\x01 \x01(\x05\x12\x16\n\x0enum_thresholds\x18\x02 \x01(\rb\x06proto3') +) + + + + +_PRCURVEPLUGINDATA = _descriptor.Descriptor( + name='PrCurvePluginData', + full_name='tensorboard.PrCurvePluginData', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='version', full_name='tensorboard.PrCurvePluginData.version', index=0, + number=1, type=5, cpp_type=1, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='num_thresholds', full_name='tensorboard.PrCurvePluginData.num_thresholds', index=1, + number=2, type=13, cpp_type=3, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=55, + serialized_end=115, +) + +DESCRIPTOR.message_types_by_name['PrCurvePluginData'] = _PRCURVEPLUGINDATA +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +PrCurvePluginData = _reflection.GeneratedProtocolMessageType('PrCurvePluginData', (_message.Message,), dict( + DESCRIPTOR = _PRCURVEPLUGINDATA, + __module__ = 'tensorboardX.src.plugin_pr_curve_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.PrCurvePluginData) + )) +_sym_db.RegisterMessage(PrCurvePluginData) + + +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/resource_handle_pb2.py b/tensorboardX/src/resource_handle_pb2.py new file mode 100644 index 0000000..276703e --- /dev/null +++ b/tensorboardX/src/resource_handle_pb2.py @@ -0,0 +1,99 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/resource_handle.proto + +import sys +_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +from google.protobuf import descriptor_pb2 +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor.FileDescriptor( + name='tensorboardX/src/resource_handle.proto', + package='tensorboard', + syntax='proto3', + serialized_pb=_b('\n&tensorboardX/src/resource_handle.proto\x12\x0btensorboard\"r\n\x13ResourceHandleProto\x12\x0e\n\x06\x64\x65vice\x18\x01 \x01(\t\x12\x11\n\tcontainer\x18\x02 \x01(\t\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x11\n\thash_code\x18\x04 \x01(\x04\x12\x17\n\x0fmaybe_type_name\x18\x05 \x01(\tB/\n\x18org.tensorflow.frameworkB\x0eResourceHandleP\x01\xf8\x01\x01\x62\x06proto3') +) + + + + +_RESOURCEHANDLEPROTO = _descriptor.Descriptor( + name='ResourceHandleProto', + full_name='tensorboard.ResourceHandleProto', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='device', full_name='tensorboard.ResourceHandleProto.device', index=0, + number=1, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='container', full_name='tensorboard.ResourceHandleProto.container', index=1, + number=2, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='name', full_name='tensorboard.ResourceHandleProto.name', index=2, + number=3, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='hash_code', full_name='tensorboard.ResourceHandleProto.hash_code', index=3, + number=4, type=4, cpp_type=4, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='maybe_type_name', full_name='tensorboard.ResourceHandleProto.maybe_type_name', index=4, + number=5, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=55, + serialized_end=169, +) + +DESCRIPTOR.message_types_by_name['ResourceHandleProto'] = _RESOURCEHANDLEPROTO +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +ResourceHandleProto = _reflection.GeneratedProtocolMessageType('ResourceHandleProto', (_message.Message,), dict( + DESCRIPTOR = _RESOURCEHANDLEPROTO, + __module__ = 'tensorboardX.src.resource_handle_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.ResourceHandleProto) + )) +_sym_db.RegisterMessage(ResourceHandleProto) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\016ResourceHandleP\001\370\001\001')) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/summary_pb2.py b/tensorboardX/src/summary_pb2.py new file mode 100644 index 0000000..db24ea8 --- /dev/null +++ b/tensorboardX/src/summary_pb2.py @@ -0,0 +1,532 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/summary.proto + +import sys +_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +from google.protobuf import descriptor_pb2 +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from tensorboardX.src import tensor_pb2 as tensorboardX_dot_src_dot_tensor__pb2 + + +DESCRIPTOR = _descriptor.FileDescriptor( + name='tensorboardX/src/summary.proto', + package='tensorboard', + syntax='proto3', + serialized_pb=_b('\n\x1etensorboardX/src/summary.proto\x12\x0btensorboard\x1a\x1dtensorboardX/src/tensor.proto\"\'\n\x12SummaryDescription\x12\x11\n\ttype_hint\x18\x01 \x01(\t\"\x87\x01\n\x0eHistogramProto\x12\x0b\n\x03min\x18\x01 \x01(\x01\x12\x0b\n\x03max\x18\x02 \x01(\x01\x12\x0b\n\x03num\x18\x03 \x01(\x01\x12\x0b\n\x03sum\x18\x04 \x01(\x01\x12\x13\n\x0bsum_squares\x18\x05 \x01(\x01\x12\x18\n\x0c\x62ucket_limit\x18\x06 \x03(\x01\x42\x02\x10\x01\x12\x12\n\x06\x62ucket\x18\x07 \x03(\x01\x42\x02\x10\x01\"\x83\x01\n\x0fSummaryMetadata\x12<\n\x0bplugin_data\x18\x01 \x03(\x0b\x32\'.tensorboard.SummaryMetadata.PluginData\x1a\x32\n\nPluginData\x12\x13\n\x0bplugin_name\x18\x01 \x01(\t\x12\x0f\n\x07\x63ontent\x18\x02 \x01(\t\"\xe4\x04\n\x07Summary\x12)\n\x05value\x18\x01 \x03(\x0b\x32\x1a.tensorboard.Summary.Value\x1aX\n\x05Image\x12\x0e\n\x06height\x18\x01 \x01(\x05\x12\r\n\x05width\x18\x02 \x01(\x05\x12\x12\n\ncolorspace\x18\x03 \x01(\x05\x12\x1c\n\x14\x65ncoded_image_string\x18\x04 \x01(\x0c\x1a}\n\x05\x41udio\x12\x13\n\x0bsample_rate\x18\x01 \x01(\x02\x12\x14\n\x0cnum_channels\x18\x02 \x01(\x03\x12\x15\n\rlength_frames\x18\x03 \x01(\x03\x12\x1c\n\x14\x65ncoded_audio_string\x18\x04 \x01(\x0c\x12\x14\n\x0c\x63ontent_type\x18\x05 \x01(\t\x1a\xd4\x02\n\x05Value\x12\x11\n\tnode_name\x18\x07 \x01(\t\x12\x0b\n\x03tag\x18\x01 \x01(\t\x12.\n\x08metadata\x18\t \x01(\x0b\x32\x1c.tensorboard.SummaryMetadata\x12\x16\n\x0csimple_value\x18\x02 \x01(\x02H\x00\x12&\n\x1cobsolete_old_style_histogram\x18\x03 \x01(\x0cH\x00\x12+\n\x05image\x18\x04 \x01(\x0b\x32\x1a.tensorboard.Summary.ImageH\x00\x12,\n\x05histo\x18\x05 \x01(\x0b\x32\x1b.tensorboard.HistogramProtoH\x00\x12+\n\x05\x61udio\x18\x06 \x01(\x0b\x32\x1a.tensorboard.Summary.AudioH\x00\x12*\n\x06tensor\x18\x08 \x01(\x0b\x32\x18.tensorboard.TensorProtoH\x00\x42\x07\n\x05valueB.\n\x18org.tensorflow.frameworkB\rSummaryProtosP\x01\xf8\x01\x01\x62\x06proto3') + , + dependencies=[tensorboardX_dot_src_dot_tensor__pb2.DESCRIPTOR,]) + + + + +_SUMMARYDESCRIPTION = _descriptor.Descriptor( + name='SummaryDescription', + full_name='tensorboard.SummaryDescription', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='type_hint', full_name='tensorboard.SummaryDescription.type_hint', index=0, + number=1, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=78, + serialized_end=117, +) + + +_HISTOGRAMPROTO = _descriptor.Descriptor( + name='HistogramProto', + full_name='tensorboard.HistogramProto', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='min', full_name='tensorboard.HistogramProto.min', index=0, + number=1, type=1, cpp_type=5, label=1, + has_default_value=False, default_value=float(0), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='max', full_name='tensorboard.HistogramProto.max', index=1, + number=2, type=1, cpp_type=5, label=1, + has_default_value=False, default_value=float(0), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='num', full_name='tensorboard.HistogramProto.num', index=2, + number=3, type=1, cpp_type=5, label=1, + has_default_value=False, default_value=float(0), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='sum', full_name='tensorboard.HistogramProto.sum', index=3, + number=4, type=1, cpp_type=5, label=1, + has_default_value=False, default_value=float(0), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='sum_squares', full_name='tensorboard.HistogramProto.sum_squares', index=4, + number=5, type=1, cpp_type=5, label=1, + has_default_value=False, default_value=float(0), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='bucket_limit', full_name='tensorboard.HistogramProto.bucket_limit', index=5, + number=6, type=1, cpp_type=5, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))), + _descriptor.FieldDescriptor( + name='bucket', full_name='tensorboard.HistogramProto.bucket', index=6, + number=7, type=1, cpp_type=5, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=120, + serialized_end=255, +) + + +_SUMMARYMETADATA_PLUGINDATA = _descriptor.Descriptor( + name='PluginData', + full_name='tensorboard.SummaryMetadata.PluginData', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='plugin_name', full_name='tensorboard.SummaryMetadata.PluginData.plugin_name', index=0, + number=1, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='content', full_name='tensorboard.SummaryMetadata.PluginData.content', index=1, + number=2, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=339, + serialized_end=389, +) + +_SUMMARYMETADATA = _descriptor.Descriptor( + name='SummaryMetadata', + full_name='tensorboard.SummaryMetadata', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='plugin_data', full_name='tensorboard.SummaryMetadata.plugin_data', index=0, + number=1, type=11, cpp_type=10, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[_SUMMARYMETADATA_PLUGINDATA, ], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=258, + serialized_end=389, +) + + +_SUMMARY_IMAGE = _descriptor.Descriptor( + name='Image', + full_name='tensorboard.Summary.Image', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='height', full_name='tensorboard.Summary.Image.height', index=0, + number=1, type=5, cpp_type=1, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='width', full_name='tensorboard.Summary.Image.width', index=1, + number=2, type=5, cpp_type=1, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='colorspace', full_name='tensorboard.Summary.Image.colorspace', index=2, + number=3, type=5, cpp_type=1, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='encoded_image_string', full_name='tensorboard.Summary.Image.encoded_image_string', index=3, + number=4, type=12, cpp_type=9, label=1, + has_default_value=False, default_value=_b(""), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=446, + serialized_end=534, +) + +_SUMMARY_AUDIO = _descriptor.Descriptor( + name='Audio', + full_name='tensorboard.Summary.Audio', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='sample_rate', full_name='tensorboard.Summary.Audio.sample_rate', index=0, + number=1, type=2, cpp_type=6, label=1, + has_default_value=False, default_value=float(0), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='num_channels', full_name='tensorboard.Summary.Audio.num_channels', index=1, + number=2, type=3, cpp_type=2, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='length_frames', full_name='tensorboard.Summary.Audio.length_frames', index=2, + number=3, type=3, cpp_type=2, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='encoded_audio_string', full_name='tensorboard.Summary.Audio.encoded_audio_string', index=3, + number=4, type=12, cpp_type=9, label=1, + has_default_value=False, default_value=_b(""), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='content_type', full_name='tensorboard.Summary.Audio.content_type', index=4, + number=5, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=536, + serialized_end=661, +) + +_SUMMARY_VALUE = _descriptor.Descriptor( + name='Value', + full_name='tensorboard.Summary.Value', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='node_name', full_name='tensorboard.Summary.Value.node_name', index=0, + number=7, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='tag', full_name='tensorboard.Summary.Value.tag', index=1, + number=1, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='metadata', full_name='tensorboard.Summary.Value.metadata', index=2, + number=9, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='simple_value', full_name='tensorboard.Summary.Value.simple_value', index=3, + number=2, type=2, cpp_type=6, label=1, + has_default_value=False, default_value=float(0), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='obsolete_old_style_histogram', full_name='tensorboard.Summary.Value.obsolete_old_style_histogram', index=4, + number=3, type=12, cpp_type=9, label=1, + has_default_value=False, default_value=_b(""), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='image', full_name='tensorboard.Summary.Value.image', index=5, + number=4, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='histo', full_name='tensorboard.Summary.Value.histo', index=6, + number=5, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='audio', full_name='tensorboard.Summary.Value.audio', index=7, + number=6, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='tensor', full_name='tensorboard.Summary.Value.tensor', index=8, + number=8, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + _descriptor.OneofDescriptor( + name='value', full_name='tensorboard.Summary.Value.value', + index=0, containing_type=None, fields=[]), + ], + serialized_start=664, + serialized_end=1004, +) + +_SUMMARY = _descriptor.Descriptor( + name='Summary', + full_name='tensorboard.Summary', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='value', full_name='tensorboard.Summary.value', index=0, + number=1, type=11, cpp_type=10, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[_SUMMARY_IMAGE, _SUMMARY_AUDIO, _SUMMARY_VALUE, ], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=392, + serialized_end=1004, +) + +_SUMMARYMETADATA_PLUGINDATA.containing_type = _SUMMARYMETADATA +_SUMMARYMETADATA.fields_by_name['plugin_data'].message_type = _SUMMARYMETADATA_PLUGINDATA +_SUMMARY_IMAGE.containing_type = _SUMMARY +_SUMMARY_AUDIO.containing_type = _SUMMARY +_SUMMARY_VALUE.fields_by_name['metadata'].message_type = _SUMMARYMETADATA +_SUMMARY_VALUE.fields_by_name['image'].message_type = _SUMMARY_IMAGE +_SUMMARY_VALUE.fields_by_name['histo'].message_type = _HISTOGRAMPROTO +_SUMMARY_VALUE.fields_by_name['audio'].message_type = _SUMMARY_AUDIO +_SUMMARY_VALUE.fields_by_name['tensor'].message_type = tensorboardX_dot_src_dot_tensor__pb2._TENSORPROTO +_SUMMARY_VALUE.containing_type = _SUMMARY +_SUMMARY_VALUE.oneofs_by_name['value'].fields.append( + _SUMMARY_VALUE.fields_by_name['simple_value']) +_SUMMARY_VALUE.fields_by_name['simple_value'].containing_oneof = _SUMMARY_VALUE.oneofs_by_name['value'] +_SUMMARY_VALUE.oneofs_by_name['value'].fields.append( + _SUMMARY_VALUE.fields_by_name['obsolete_old_style_histogram']) +_SUMMARY_VALUE.fields_by_name['obsolete_old_style_histogram'].containing_oneof = _SUMMARY_VALUE.oneofs_by_name['value'] +_SUMMARY_VALUE.oneofs_by_name['value'].fields.append( + _SUMMARY_VALUE.fields_by_name['image']) +_SUMMARY_VALUE.fields_by_name['image'].containing_oneof = _SUMMARY_VALUE.oneofs_by_name['value'] +_SUMMARY_VALUE.oneofs_by_name['value'].fields.append( + _SUMMARY_VALUE.fields_by_name['histo']) +_SUMMARY_VALUE.fields_by_name['histo'].containing_oneof = _SUMMARY_VALUE.oneofs_by_name['value'] +_SUMMARY_VALUE.oneofs_by_name['value'].fields.append( + _SUMMARY_VALUE.fields_by_name['audio']) +_SUMMARY_VALUE.fields_by_name['audio'].containing_oneof = _SUMMARY_VALUE.oneofs_by_name['value'] +_SUMMARY_VALUE.oneofs_by_name['value'].fields.append( + _SUMMARY_VALUE.fields_by_name['tensor']) +_SUMMARY_VALUE.fields_by_name['tensor'].containing_oneof = _SUMMARY_VALUE.oneofs_by_name['value'] +_SUMMARY.fields_by_name['value'].message_type = _SUMMARY_VALUE +DESCRIPTOR.message_types_by_name['SummaryDescription'] = _SUMMARYDESCRIPTION +DESCRIPTOR.message_types_by_name['HistogramProto'] = _HISTOGRAMPROTO +DESCRIPTOR.message_types_by_name['SummaryMetadata'] = _SUMMARYMETADATA +DESCRIPTOR.message_types_by_name['Summary'] = _SUMMARY +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +SummaryDescription = _reflection.GeneratedProtocolMessageType('SummaryDescription', (_message.Message,), dict( + DESCRIPTOR = _SUMMARYDESCRIPTION, + __module__ = 'tensorboardX.src.summary_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.SummaryDescription) + )) +_sym_db.RegisterMessage(SummaryDescription) + +HistogramProto = _reflection.GeneratedProtocolMessageType('HistogramProto', (_message.Message,), dict( + DESCRIPTOR = _HISTOGRAMPROTO, + __module__ = 'tensorboardX.src.summary_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.HistogramProto) + )) +_sym_db.RegisterMessage(HistogramProto) + +SummaryMetadata = _reflection.GeneratedProtocolMessageType('SummaryMetadata', (_message.Message,), dict( + + PluginData = _reflection.GeneratedProtocolMessageType('PluginData', (_message.Message,), dict( + DESCRIPTOR = _SUMMARYMETADATA_PLUGINDATA, + __module__ = 'tensorboardX.src.summary_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.SummaryMetadata.PluginData) + )) + , + DESCRIPTOR = _SUMMARYMETADATA, + __module__ = 'tensorboardX.src.summary_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.SummaryMetadata) + )) +_sym_db.RegisterMessage(SummaryMetadata) +_sym_db.RegisterMessage(SummaryMetadata.PluginData) + +Summary = _reflection.GeneratedProtocolMessageType('Summary', (_message.Message,), dict( + + Image = _reflection.GeneratedProtocolMessageType('Image', (_message.Message,), dict( + DESCRIPTOR = _SUMMARY_IMAGE, + __module__ = 'tensorboardX.src.summary_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.Summary.Image) + )) + , + + Audio = _reflection.GeneratedProtocolMessageType('Audio', (_message.Message,), dict( + DESCRIPTOR = _SUMMARY_AUDIO, + __module__ = 'tensorboardX.src.summary_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.Summary.Audio) + )) + , + + Value = _reflection.GeneratedProtocolMessageType('Value', (_message.Message,), dict( + DESCRIPTOR = _SUMMARY_VALUE, + __module__ = 'tensorboardX.src.summary_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.Summary.Value) + )) + , + DESCRIPTOR = _SUMMARY, + __module__ = 'tensorboardX.src.summary_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.Summary) + )) +_sym_db.RegisterMessage(Summary) +_sym_db.RegisterMessage(Summary.Image) +_sym_db.RegisterMessage(Summary.Audio) +_sym_db.RegisterMessage(Summary.Value) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\rSummaryProtosP\001\370\001\001')) +_HISTOGRAMPROTO.fields_by_name['bucket_limit'].has_options = True +_HISTOGRAMPROTO.fields_by_name['bucket_limit']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001')) +_HISTOGRAMPROTO.fields_by_name['bucket'].has_options = True +_HISTOGRAMPROTO.fields_by_name['bucket']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001')) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/tensor_pb2.py b/tensorboardX/src/tensor_pb2.py new file mode 100644 index 0000000..b54d9c1 --- /dev/null +++ b/tensorboardX/src/tensor_pb2.py @@ -0,0 +1,185 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/tensor.proto + +import sys +_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +from google.protobuf import descriptor_pb2 +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from tensorboardX.src import resource_handle_pb2 as tensorboardX_dot_src_dot_resource__handle__pb2 +from tensorboardX.src import tensor_shape_pb2 as tensorboardX_dot_src_dot_tensor__shape__pb2 +from tensorboardX.src import types_pb2 as tensorboardX_dot_src_dot_types__pb2 + + +DESCRIPTOR = _descriptor.FileDescriptor( + name='tensorboardX/src/tensor.proto', + package='tensorboard', + syntax='proto3', + serialized_pb=_b('\n\x1dtensorboardX/src/tensor.proto\x12\x0btensorboard\x1a&tensorboardX/src/resource_handle.proto\x1a#tensorboardX/src/tensor_shape.proto\x1a\x1ctensorboardX/src/types.proto\"\xa6\x03\n\x0bTensorProto\x12$\n\x05\x64type\x18\x01 \x01(\x0e\x32\x15.tensorboard.DataType\x12\x33\n\x0ctensor_shape\x18\x02 \x01(\x0b\x32\x1d.tensorboard.TensorShapeProto\x12\x16\n\x0eversion_number\x18\x03 \x01(\x05\x12\x16\n\x0etensor_content\x18\x04 \x01(\x0c\x12\x14\n\x08half_val\x18\r \x03(\x05\x42\x02\x10\x01\x12\x15\n\tfloat_val\x18\x05 \x03(\x02\x42\x02\x10\x01\x12\x16\n\ndouble_val\x18\x06 \x03(\x01\x42\x02\x10\x01\x12\x13\n\x07int_val\x18\x07 \x03(\x05\x42\x02\x10\x01\x12\x12\n\nstring_val\x18\x08 \x03(\x0c\x12\x18\n\x0cscomplex_val\x18\t \x03(\x02\x42\x02\x10\x01\x12\x15\n\tint64_val\x18\n \x03(\x03\x42\x02\x10\x01\x12\x14\n\x08\x62ool_val\x18\x0b \x03(\x08\x42\x02\x10\x01\x12\x18\n\x0c\x64\x63omplex_val\x18\x0c \x03(\x01\x42\x02\x10\x01\x12=\n\x13resource_handle_val\x18\x0e \x03(\x0b\x32 .tensorboard.ResourceHandleProtoB-\n\x18org.tensorflow.frameworkB\x0cTensorProtosP\x01\xf8\x01\x01\x62\x06proto3') + , + dependencies=[tensorboardX_dot_src_dot_resource__handle__pb2.DESCRIPTOR,tensorboardX_dot_src_dot_tensor__shape__pb2.DESCRIPTOR,tensorboardX_dot_src_dot_types__pb2.DESCRIPTOR,]) + + + + +_TENSORPROTO = _descriptor.Descriptor( + name='TensorProto', + full_name='tensorboard.TensorProto', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='dtype', full_name='tensorboard.TensorProto.dtype', index=0, + number=1, type=14, cpp_type=8, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='tensor_shape', full_name='tensorboard.TensorProto.tensor_shape', index=1, + number=2, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='version_number', full_name='tensorboard.TensorProto.version_number', index=2, + number=3, type=5, cpp_type=1, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='tensor_content', full_name='tensorboard.TensorProto.tensor_content', index=3, + number=4, type=12, cpp_type=9, label=1, + has_default_value=False, default_value=_b(""), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='half_val', full_name='tensorboard.TensorProto.half_val', index=4, + number=13, type=5, cpp_type=1, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))), + _descriptor.FieldDescriptor( + name='float_val', full_name='tensorboard.TensorProto.float_val', index=5, + number=5, type=2, cpp_type=6, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))), + _descriptor.FieldDescriptor( + name='double_val', full_name='tensorboard.TensorProto.double_val', index=6, + number=6, type=1, cpp_type=5, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))), + _descriptor.FieldDescriptor( + name='int_val', full_name='tensorboard.TensorProto.int_val', index=7, + number=7, type=5, cpp_type=1, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))), + _descriptor.FieldDescriptor( + name='string_val', full_name='tensorboard.TensorProto.string_val', index=8, + number=8, type=12, cpp_type=9, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='scomplex_val', full_name='tensorboard.TensorProto.scomplex_val', index=9, + number=9, type=2, cpp_type=6, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))), + _descriptor.FieldDescriptor( + name='int64_val', full_name='tensorboard.TensorProto.int64_val', index=10, + number=10, type=3, cpp_type=2, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))), + _descriptor.FieldDescriptor( + name='bool_val', full_name='tensorboard.TensorProto.bool_val', index=11, + number=11, type=8, cpp_type=7, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))), + _descriptor.FieldDescriptor( + name='dcomplex_val', full_name='tensorboard.TensorProto.dcomplex_val', index=12, + number=12, type=1, cpp_type=5, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))), + _descriptor.FieldDescriptor( + name='resource_handle_val', full_name='tensorboard.TensorProto.resource_handle_val', index=13, + number=14, type=11, cpp_type=10, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=154, + serialized_end=576, +) + +_TENSORPROTO.fields_by_name['dtype'].enum_type = tensorboardX_dot_src_dot_types__pb2._DATATYPE +_TENSORPROTO.fields_by_name['tensor_shape'].message_type = tensorboardX_dot_src_dot_tensor__shape__pb2._TENSORSHAPEPROTO +_TENSORPROTO.fields_by_name['resource_handle_val'].message_type = tensorboardX_dot_src_dot_resource__handle__pb2._RESOURCEHANDLEPROTO +DESCRIPTOR.message_types_by_name['TensorProto'] = _TENSORPROTO +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +TensorProto = _reflection.GeneratedProtocolMessageType('TensorProto', (_message.Message,), dict( + DESCRIPTOR = _TENSORPROTO, + __module__ = 'tensorboardX.src.tensor_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.TensorProto) + )) +_sym_db.RegisterMessage(TensorProto) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\014TensorProtosP\001\370\001\001')) +_TENSORPROTO.fields_by_name['half_val'].has_options = True +_TENSORPROTO.fields_by_name['half_val']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001')) +_TENSORPROTO.fields_by_name['float_val'].has_options = True +_TENSORPROTO.fields_by_name['float_val']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001')) +_TENSORPROTO.fields_by_name['double_val'].has_options = True +_TENSORPROTO.fields_by_name['double_val']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001')) +_TENSORPROTO.fields_by_name['int_val'].has_options = True +_TENSORPROTO.fields_by_name['int_val']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001')) +_TENSORPROTO.fields_by_name['scomplex_val'].has_options = True +_TENSORPROTO.fields_by_name['scomplex_val']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001')) +_TENSORPROTO.fields_by_name['int64_val'].has_options = True +_TENSORPROTO.fields_by_name['int64_val']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001')) +_TENSORPROTO.fields_by_name['bool_val'].has_options = True +_TENSORPROTO.fields_by_name['bool_val']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001')) +_TENSORPROTO.fields_by_name['dcomplex_val'].has_options = True +_TENSORPROTO.fields_by_name['dcomplex_val']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001')) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/tensor_shape_pb2.py b/tensorboardX/src/tensor_shape_pb2.py new file mode 100644 index 0000000..4934cb6 --- /dev/null +++ b/tensorboardX/src/tensor_shape_pb2.py @@ -0,0 +1,125 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/tensor_shape.proto + +import sys +_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +from google.protobuf import descriptor_pb2 +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor.FileDescriptor( + name='tensorboardX/src/tensor_shape.proto', + package='tensorboard', + syntax='proto3', + serialized_pb=_b('\n#tensorboardX/src/tensor_shape.proto\x12\x0btensorboard\"{\n\x10TensorShapeProto\x12.\n\x03\x64im\x18\x02 \x03(\x0b\x32!.tensorboard.TensorShapeProto.Dim\x12\x14\n\x0cunknown_rank\x18\x03 \x01(\x08\x1a!\n\x03\x44im\x12\x0c\n\x04size\x18\x01 \x01(\x03\x12\x0c\n\x04name\x18\x02 \x01(\tB2\n\x18org.tensorflow.frameworkB\x11TensorShapeProtosP\x01\xf8\x01\x01\x62\x06proto3') +) + + + + +_TENSORSHAPEPROTO_DIM = _descriptor.Descriptor( + name='Dim', + full_name='tensorboard.TensorShapeProto.Dim', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='size', full_name='tensorboard.TensorShapeProto.Dim.size', index=0, + number=1, type=3, cpp_type=2, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='name', full_name='tensorboard.TensorShapeProto.Dim.name', index=1, + number=2, type=9, cpp_type=9, label=1, + has_default_value=False, default_value=_b("").decode('utf-8'), + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=142, + serialized_end=175, +) + +_TENSORSHAPEPROTO = _descriptor.Descriptor( + name='TensorShapeProto', + full_name='tensorboard.TensorShapeProto', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='dim', full_name='tensorboard.TensorShapeProto.dim', index=0, + number=2, type=11, cpp_type=10, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='unknown_rank', full_name='tensorboard.TensorShapeProto.unknown_rank', index=1, + number=3, type=8, cpp_type=7, label=1, + has_default_value=False, default_value=False, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[_TENSORSHAPEPROTO_DIM, ], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=52, + serialized_end=175, +) + +_TENSORSHAPEPROTO_DIM.containing_type = _TENSORSHAPEPROTO +_TENSORSHAPEPROTO.fields_by_name['dim'].message_type = _TENSORSHAPEPROTO_DIM +DESCRIPTOR.message_types_by_name['TensorShapeProto'] = _TENSORSHAPEPROTO +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +TensorShapeProto = _reflection.GeneratedProtocolMessageType('TensorShapeProto', (_message.Message,), dict( + + Dim = _reflection.GeneratedProtocolMessageType('Dim', (_message.Message,), dict( + DESCRIPTOR = _TENSORSHAPEPROTO_DIM, + __module__ = 'tensorboardX.src.tensor_shape_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.TensorShapeProto.Dim) + )) + , + DESCRIPTOR = _TENSORSHAPEPROTO, + __module__ = 'tensorboardX.src.tensor_shape_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.TensorShapeProto) + )) +_sym_db.RegisterMessage(TensorShapeProto) +_sym_db.RegisterMessage(TensorShapeProto.Dim) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\021TensorShapeProtosP\001\370\001\001')) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/types_pb2.py b/tensorboardX/src/types_pb2.py new file mode 100644 index 0000000..cd02d02 --- /dev/null +++ b/tensorboardX/src/types_pb2.py @@ -0,0 +1,254 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/types.proto + +import sys +_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) +from google.protobuf.internal import enum_type_wrapper +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +from google.protobuf import descriptor_pb2 +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor.FileDescriptor( + name='tensorboardX/src/types.proto', + package='tensorboard', + syntax='proto3', + serialized_pb=_b('\n\x1ctensorboardX/src/types.proto\x12\x0btensorboard*\xc2\x05\n\x08\x44\x61taType\x12\x0e\n\nDT_INVALID\x10\x00\x12\x0c\n\x08\x44T_FLOAT\x10\x01\x12\r\n\tDT_DOUBLE\x10\x02\x12\x0c\n\x08\x44T_INT32\x10\x03\x12\x0c\n\x08\x44T_UINT8\x10\x04\x12\x0c\n\x08\x44T_INT16\x10\x05\x12\x0b\n\x07\x44T_INT8\x10\x06\x12\r\n\tDT_STRING\x10\x07\x12\x10\n\x0c\x44T_COMPLEX64\x10\x08\x12\x0c\n\x08\x44T_INT64\x10\t\x12\x0b\n\x07\x44T_BOOL\x10\n\x12\x0c\n\x08\x44T_QINT8\x10\x0b\x12\r\n\tDT_QUINT8\x10\x0c\x12\r\n\tDT_QINT32\x10\r\x12\x0f\n\x0b\x44T_BFLOAT16\x10\x0e\x12\r\n\tDT_QINT16\x10\x0f\x12\x0e\n\nDT_QUINT16\x10\x10\x12\r\n\tDT_UINT16\x10\x11\x12\x11\n\rDT_COMPLEX128\x10\x12\x12\x0b\n\x07\x44T_HALF\x10\x13\x12\x0f\n\x0b\x44T_RESOURCE\x10\x14\x12\x10\n\x0c\x44T_FLOAT_REF\x10\x65\x12\x11\n\rDT_DOUBLE_REF\x10\x66\x12\x10\n\x0c\x44T_INT32_REF\x10g\x12\x10\n\x0c\x44T_UINT8_REF\x10h\x12\x10\n\x0c\x44T_INT16_REF\x10i\x12\x0f\n\x0b\x44T_INT8_REF\x10j\x12\x11\n\rDT_STRING_REF\x10k\x12\x14\n\x10\x44T_COMPLEX64_REF\x10l\x12\x10\n\x0c\x44T_INT64_REF\x10m\x12\x0f\n\x0b\x44T_BOOL_REF\x10n\x12\x10\n\x0c\x44T_QINT8_REF\x10o\x12\x11\n\rDT_QUINT8_REF\x10p\x12\x11\n\rDT_QINT32_REF\x10q\x12\x13\n\x0f\x44T_BFLOAT16_REF\x10r\x12\x11\n\rDT_QINT16_REF\x10s\x12\x12\n\x0e\x44T_QUINT16_REF\x10t\x12\x11\n\rDT_UINT16_REF\x10u\x12\x15\n\x11\x44T_COMPLEX128_REF\x10v\x12\x0f\n\x0b\x44T_HALF_REF\x10w\x12\x13\n\x0f\x44T_RESOURCE_REF\x10xB,\n\x18org.tensorflow.frameworkB\x0bTypesProtosP\x01\xf8\x01\x01\x62\x06proto3') +) + +_DATATYPE = _descriptor.EnumDescriptor( + name='DataType', + full_name='tensorboard.DataType', + filename=None, + file=DESCRIPTOR, + values=[ + _descriptor.EnumValueDescriptor( + name='DT_INVALID', index=0, number=0, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_FLOAT', index=1, number=1, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_DOUBLE', index=2, number=2, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_INT32', index=3, number=3, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_UINT8', index=4, number=4, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_INT16', index=5, number=5, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_INT8', index=6, number=6, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_STRING', index=7, number=7, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_COMPLEX64', index=8, number=8, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_INT64', index=9, number=9, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_BOOL', index=10, number=10, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_QINT8', index=11, number=11, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_QUINT8', index=12, number=12, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_QINT32', index=13, number=13, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_BFLOAT16', index=14, number=14, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_QINT16', index=15, number=15, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_QUINT16', index=16, number=16, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_UINT16', index=17, number=17, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_COMPLEX128', index=18, number=18, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_HALF', index=19, number=19, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_RESOURCE', index=20, number=20, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_FLOAT_REF', index=21, number=101, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_DOUBLE_REF', index=22, number=102, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_INT32_REF', index=23, number=103, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_UINT8_REF', index=24, number=104, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_INT16_REF', index=25, number=105, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_INT8_REF', index=26, number=106, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_STRING_REF', index=27, number=107, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_COMPLEX64_REF', index=28, number=108, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_INT64_REF', index=29, number=109, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_BOOL_REF', index=30, number=110, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_QINT8_REF', index=31, number=111, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_QUINT8_REF', index=32, number=112, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_QINT32_REF', index=33, number=113, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_BFLOAT16_REF', index=34, number=114, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_QINT16_REF', index=35, number=115, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_QUINT16_REF', index=36, number=116, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_UINT16_REF', index=37, number=117, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_COMPLEX128_REF', index=38, number=118, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_HALF_REF', index=39, number=119, + options=None, + type=None), + _descriptor.EnumValueDescriptor( + name='DT_RESOURCE_REF', index=40, number=120, + options=None, + type=None), + ], + containing_type=None, + options=None, + serialized_start=46, + serialized_end=752, +) +_sym_db.RegisterEnumDescriptor(_DATATYPE) + +DataType = enum_type_wrapper.EnumTypeWrapper(_DATATYPE) +DT_INVALID = 0 +DT_FLOAT = 1 +DT_DOUBLE = 2 +DT_INT32 = 3 +DT_UINT8 = 4 +DT_INT16 = 5 +DT_INT8 = 6 +DT_STRING = 7 +DT_COMPLEX64 = 8 +DT_INT64 = 9 +DT_BOOL = 10 +DT_QINT8 = 11 +DT_QUINT8 = 12 +DT_QINT32 = 13 +DT_BFLOAT16 = 14 +DT_QINT16 = 15 +DT_QUINT16 = 16 +DT_UINT16 = 17 +DT_COMPLEX128 = 18 +DT_HALF = 19 +DT_RESOURCE = 20 +DT_FLOAT_REF = 101 +DT_DOUBLE_REF = 102 +DT_INT32_REF = 103 +DT_UINT8_REF = 104 +DT_INT16_REF = 105 +DT_INT8_REF = 106 +DT_STRING_REF = 107 +DT_COMPLEX64_REF = 108 +DT_INT64_REF = 109 +DT_BOOL_REF = 110 +DT_QINT8_REF = 111 +DT_QUINT8_REF = 112 +DT_QINT32_REF = 113 +DT_BFLOAT16_REF = 114 +DT_QINT16_REF = 115 +DT_QUINT16_REF = 116 +DT_UINT16_REF = 117 +DT_COMPLEX128_REF = 118 +DT_HALF_REF = 119 +DT_RESOURCE_REF = 120 + + +DESCRIPTOR.enum_types_by_name['DataType'] = _DATATYPE +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\013TypesProtosP\001\370\001\001')) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/versions_pb2.py b/tensorboardX/src/versions_pb2.py new file mode 100644 index 0000000..1daceb7 --- /dev/null +++ b/tensorboardX/src/versions_pb2.py @@ -0,0 +1,85 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/versions.proto + +import sys +_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +from google.protobuf import descriptor_pb2 +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor.FileDescriptor( + name='tensorboardX/src/versions.proto', + package='tensorboard', + syntax='proto3', + serialized_pb=_b('\n\x1ftensorboardX/src/versions.proto\x12\x0btensorboard\"K\n\nVersionDef\x12\x10\n\x08producer\x18\x01 \x01(\x05\x12\x14\n\x0cmin_consumer\x18\x02 \x01(\x05\x12\x15\n\rbad_consumers\x18\x03 \x03(\x05\x42/\n\x18org.tensorflow.frameworkB\x0eVersionsProtosP\x01\xf8\x01\x01\x62\x06proto3') +) + + + + +_VERSIONDEF = _descriptor.Descriptor( + name='VersionDef', + full_name='tensorboard.VersionDef', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='producer', full_name='tensorboard.VersionDef.producer', index=0, + number=1, type=5, cpp_type=1, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='min_consumer', full_name='tensorboard.VersionDef.min_consumer', index=1, + number=2, type=5, cpp_type=1, label=1, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='bad_consumers', full_name='tensorboard.VersionDef.bad_consumers', index=2, + number=3, type=5, cpp_type=1, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[ + ], + serialized_start=48, + serialized_end=123, +) + +DESCRIPTOR.message_types_by_name['VersionDef'] = _VERSIONDEF +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +VersionDef = _reflection.GeneratedProtocolMessageType('VersionDef', (_message.Message,), dict( + DESCRIPTOR = _VERSIONDEF, + __module__ = 'tensorboardX.src.versions_pb2' + # @@protoc_insertion_point(class_scope:tensorboard.VersionDef) + )) +_sym_db.RegisterMessage(VersionDef) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\016VersionsProtosP\001\370\001\001')) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/summary.py b/tensorboardX/summary.py new file mode 100644 index 0000000..80d3c69 --- /dev/null +++ b/tensorboardX/summary.py @@ -0,0 +1,264 @@ +# Copyright 2016 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""## Generation of summaries. +### Class for writing Summaries +@@FileWriter +@@FileWriterCache +### Summary Ops +@@tensor_summary +@@scalar +@@histogram +@@audio +@@image +@@merge +@@merge_all +## Utilities +@@get_summary_description +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import logging +import re as _re +import bisect +from six import StringIO +from six.moves import range +import numpy as np +# pylint: disable=unused-import +from .src.summary_pb2 import Summary +from .src.summary_pb2 import HistogramProto +from .src.summary_pb2 import SummaryMetadata +from .src.tensor_pb2 import TensorProto +from .src.tensor_shape_pb2 import TensorShapeProto +from .src.plugin_pr_curve_pb2 import PrCurvePluginData +from .x2num import makenp + +_INVALID_TAG_CHARACTERS = _re.compile(r'[^-/\w\.]') + + +def _clean_tag(name): + # In the past, the first argument to summary ops was a tag, which allowed + # arbitrary characters. Now we are changing the first argument to be the node + # name. This has a number of advantages (users of summary ops now can + # take advantage of the tf name scope system) but risks breaking existing + # usage, because a much smaller set of characters are allowed in node names. + # This function replaces all illegal characters with _s, and logs a warning. + # It also strips leading slashes from the name. + if name is not None: + new_name = _INVALID_TAG_CHARACTERS.sub('_', name) + new_name = new_name.lstrip('/') # Remove leading slashes + if new_name != name: + logging.info('Summary name %s is illegal; using %s instead.' % (name, new_name)) + name = new_name + return name + + +def scalar(name, scalar, collections=None): + """Outputs a `Summary` protocol buffer containing a single scalar value. + The generated Summary has a Tensor.proto containing the input Tensor. + Args: + name: A name for the generated node. Will also serve as the series name in + TensorBoard. + tensor: A real numeric Tensor containing a single value. + collections: Optional list of graph collections keys. The new summary op is + added to these collections. Defaults to `[GraphKeys.SUMMARIES]`. + Returns: + A scalar `Tensor` of type `string`. Which contains a `Summary` protobuf. + Raises: + ValueError: If tensor has the wrong shape or type. + """ + name = _clean_tag(name) + scalar = makenp(scalar) + assert(scalar.squeeze().ndim == 0), 'scalar should be 0D' + scalar = float(scalar) + return Summary(value=[Summary.Value(tag=name, simple_value=scalar)]) + + +def histogram(name, values, bins, collections=None): + # pylint: disable=line-too-long + """Outputs a `Summary` protocol buffer with a histogram. + The generated + [`Summary`](https://www.tensorflow.org/code/tensorflow/core/framework/summary.proto) + has one summary value containing a histogram for `values`. + This op reports an `InvalidArgument` error if any value is not finite. + Args: + name: A name for the generated node. Will also serve as a series name in + TensorBoard. + values: A real numeric `Tensor`. Any shape. Values to use to + build the histogram. + collections: Optional list of graph collections keys. The new summary op is + added to these collections. Defaults to `[GraphKeys.SUMMARIES]`. + Returns: + A scalar `Tensor` of type `string`. The serialized `Summary` protocol + buffer. + """ + name = _clean_tag(name) + values = makenp(values) + hist = make_histogram(values.astype(float), bins) + return Summary(value=[Summary.Value(tag=name, histo=hist)]) + + +def make_histogram(values, bins): + """Convert values into a histogram proto using logic from histogram.cc.""" + values = values.reshape(-1) + counts, limits = np.histogram(values, bins=bins) + limits = limits[1:] + + sum_sq = values.dot(values) + return HistogramProto(min=values.min(), + max=values.max(), + num=len(values), + sum=values.sum(), + sum_squares=sum_sq, + bucket_limit=limits, + bucket=counts) + + +def image(tag, tensor): + """Outputs a `Summary` protocol buffer with images. + The summary has up to `max_images` summary values containing images. The + images are built from `tensor` which must be 3-D with shape `[height, width, + channels]` and where `channels` can be: + * 1: `tensor` is interpreted as Grayscale. + * 3: `tensor` is interpreted as RGB. + * 4: `tensor` is interpreted as RGBA. + The `name` in the outputted Summary.Value protobufs is generated based on the + name, with a suffix depending on the max_outputs setting: + * If `max_outputs` is 1, the summary value tag is '*name*/image'. + * If `max_outputs` is greater than 1, the summary value tags are + generated sequentially as '*name*/image/0', '*name*/image/1', etc. + Args: + tag: A name for the generated node. Will also serve as a series name in + TensorBoard. + tensor: A 3-D `uint8` or `float32` `Tensor` of shape `[height, width, + channels]` where `channels` is 1, 3, or 4. + Returns: + A scalar `Tensor` of type `string`. The serialized `Summary` protocol + buffer. + """ + tag = _clean_tag(tag) + tensor = makenp(tensor, 'IMG') + tensor = tensor.astype(np.float32) + tensor = (tensor * 255).astype(np.uint8) + image = make_image(tensor) + return Summary(value=[Summary.Value(tag=tag, image=image)]) + + +def make_image(tensor): + """Convert an numpy representation image to Image protobuf""" + from PIL import Image + height, width, channel = tensor.shape + image = Image.fromarray(tensor) + import io + output = io.BytesIO() + image.save(output, format='PNG') + image_string = output.getvalue() + output.close() + return Summary.Image(height=height, + width=width, + colorspace=channel, + encoded_image_string=image_string) + + +def audio(tag, tensor, sample_rate=44100): + tensor = makenp(tensor) + tensor = tensor.squeeze() + if abs(tensor).max() > 1: + print('warning: audio amplitude out of range, auto clipped.') + tensor = tensor.clip(-1, 1) + assert(tensor.ndim == 1), 'input tensor should be 1 dimensional.' + + tensor_list = [int(32767.0 * x) for x in tensor] + import io + import wave + import struct + fio = io.BytesIO() + Wave_write = wave.open(fio, 'wb') + Wave_write.setnchannels(1) + Wave_write.setsampwidth(2) + Wave_write.setframerate(sample_rate) + tensor_enc = b'' + for v in tensor_list: + tensor_enc += struct.pack(' 127: # wierd, value > 127 breaks protobuf + num_thresholds = 127 + data = compute_curve(labels, predictions, num_thresholds=num_thresholds, weights=weights) + pr_curve_plugin_data = PrCurvePluginData(version=0, num_thresholds=num_thresholds).SerializeToString() + PluginData = [SummaryMetadata.PluginData(plugin_name='pr_curves', content=pr_curve_plugin_data)] + smd = SummaryMetadata(plugin_data=PluginData) + tensor = TensorProto(dtype='DT_FLOAT', + float_val=data.reshape(-1).tolist(), + tensor_shape=TensorShapeProto( + dim=[TensorShapeProto.Dim(size=data.shape[0]), TensorShapeProto.Dim(size=data.shape[1])])) + return Summary(value=[Summary.Value(tag=tag, metadata=smd, tensor=tensor)]) + + +# https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/pr_curve/summary.py +def compute_curve(labels, predictions, num_thresholds=None, weights=None): + _MINIMUM_COUNT = 1e-7 + + if weights is None: + weights = 1.0 + + # Compute bins of true positives and false positives. + bucket_indices = np.int32(np.floor(predictions * (num_thresholds - 1))) + float_labels = labels.astype(np.float) + histogram_range = (0, num_thresholds - 1) + tp_buckets, _ = np.histogram( + bucket_indices, + bins=num_thresholds, + range=histogram_range, + weights=float_labels * weights) + fp_buckets, _ = np.histogram( + bucket_indices, + bins=num_thresholds, + range=histogram_range, + weights=(1.0 - float_labels) * weights) + + # Obtain the reverse cumulative sum. + tp = np.cumsum(tp_buckets[::-1])[::-1] + fp = np.cumsum(fp_buckets[::-1])[::-1] + tn = fp[0] - fp + fn = tp[0] - tp + precision = tp / np.maximum(_MINIMUM_COUNT, tp + fp) + recall = tp / np.maximum(_MINIMUM_COUNT, tp + fn) + return np.stack((tp, fp, tn, fn, precision, recall)) diff --git a/tensorboardX/writer.py b/tensorboardX/writer.py new file mode 100644 index 0000000..3578c1c --- /dev/null +++ b/tensorboardX/writer.py @@ -0,0 +1,485 @@ +# Copyright 2016 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Provides an API for generating Event protocol buffers.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import time +import json +import os +from .src import event_pb2 +from .src import summary_pb2 +from .src import graph_pb2 +from .event_file_writer import EventFileWriter +from .summary import scalar, histogram, image, audio, text, pr_curve +from .graph import graph +from .graph_onnx import gg +from .embedding import make_mat, make_sprite, make_tsv, append_pbtxt + + +class SummaryToEventTransformer(object): + """Abstractly implements the SummaryWriter API. + This API basically implements a number of endpoints (add_summary, + add_session_log, etc). The endpoints all generate an event protobuf, which is + passed to the contained event_writer. + @@__init__ + @@add_summary + @@add_session_log + @@add_graph + @@add_meta_graph + @@add_run_metadata + """ + + def __init__(self, event_writer, graph=None, graph_def=None): + """Creates a `SummaryWriter` and an event file. + On construction the summary writer creates a new event file in `logdir`. + This event file will contain `Event` protocol buffers constructed when you + call one of the following functions: `add_summary()`, `add_session_log()`, + `add_event()`, or `add_graph()`. + If you pass a `Graph` to the constructor it is added to + the event file. (This is equivalent to calling `add_graph()` later). + TensorBoard will pick the graph from the file and display it graphically so + you can interactively explore the graph you built. You will usually pass + the graph from the session in which you launched it: + ```python + ...create a graph... + # Launch the graph in a session. + sess = tf.Session() + # Create a summary writer, add the 'graph' to the event file. + writer = tf.summary.FileWriter(, sess.graph) + ``` + Args: + event_writer: An EventWriter. Implements add_event method. + graph: A `Graph` object, such as `sess.graph`. + graph_def: DEPRECATED: Use the `graph` argument instead. + """ + self.event_writer = event_writer + # For storing used tags for session.run() outputs. + self._session_run_tags = {} + # TODO(zihaolucky). pass this an empty graph to check whether it's necessary. + # currently we don't support graph in MXNet using tensorboard. + + def add_summary(self, summary, global_step=None): + """Adds a `Summary` protocol buffer to the event file. + This method wraps the provided summary in an `Event` protocol buffer + and adds it to the event file. + You can pass the result of evaluating any summary op, using + [`Session.run()`](client.md#Session.run) or + [`Tensor.eval()`](framework.md#Tensor.eval), to this + function. Alternatively, you can pass a `tf.Summary` protocol + buffer that you populate with your own data. The latter is + commonly done to report evaluation results in event files. + Args: + summary: A `Summary` protocol buffer, optionally serialized as a string. + global_step: Number. Optional global step value to record with the + summary. + """ + if isinstance(summary, bytes): + summ = summary_pb2.Summary() + summ.ParseFromString(summary) + summary = summ + event = event_pb2.Event(summary=summary) + self._add_event(event, global_step) + + def add_graph_onnx(self, graph): + """Adds a `Graph` protocol buffer to the event file. + """ + event = event_pb2.Event(graph_def=graph.SerializeToString()) + self._add_event(event, None) + + def add_graph(self, graph): + """Adds a `Graph` protocol buffer to the event file. + """ + event = event_pb2.Event(graph_def=graph.SerializeToString()) + self._add_event(event, None) + + def add_session_log(self, session_log, global_step=None): + """Adds a `SessionLog` protocol buffer to the event file. + This method wraps the provided session in an `Event` protocol buffer + and adds it to the event file. + Args: + session_log: A `SessionLog` protocol buffer. + global_step: Number. Optional global step value to record with the + summary. + """ + event = event_pb2.Event(session_log=session_log) + self._add_event(event, global_step) + + def _add_event(self, event, step): + event.wall_time = time.time() + if step is not None: + event.step = int(step) + self.event_writer.add_event(event) + + +class FileWriter(SummaryToEventTransformer): + """Writes `Summary` protocol buffers to event files. + The `FileWriter` class provides a mechanism to create an event file in a + given directory and add summaries and events to it. The class updates the + file contents asynchronously. This allows a training program to call methods + to add data to the file directly from the training loop, without slowing down + training. + @@__init__ + @@add_summary + @@add_session_log + @@add_event + @@add_graph + @@add_run_metadata + @@get_logdir + @@flush + @@close + """ + + def __init__(self, + logdir, + graph=None, + max_queue=10, + flush_secs=120, + graph_def=None): + """Creates a `FileWriter` and an event file. + On construction the summary writer creates a new event file in `logdir`. + This event file will contain `Event` protocol buffers constructed when you + call one of the following functions: `add_summary()`, `add_session_log()`, + `add_event()`, or `add_graph()`. + If you pass a `Graph` to the constructor it is added to + the event file. (This is equivalent to calling `add_graph()` later). + TensorBoard will pick the graph from the file and display it graphically so + you can interactively explore the graph you built. You will usually pass + the graph from the session in which you launched it: + ```python + ...create a graph... + # Launch the graph in a session. + sess = tf.Session() + # Create a summary writer, add the 'graph' to the event file. + writer = tf.summary.FileWriter(, sess.graph) + ``` + The other arguments to the constructor control the asynchronous writes to + the event file: + * `flush_secs`: How often, in seconds, to flush the added summaries + and events to disk. + * `max_queue`: Maximum number of summaries or events pending to be + written to disk before one of the 'add' calls block. + Args: + logdir: A string. Directory where event file will be written. + graph: A `Graph` object, such as `sess.graph`. + max_queue: Integer. Size of the queue for pending events and summaries. + flush_secs: Number. How often, in seconds, to flush the + pending events and summaries to disk. + graph_def: DEPRECATED: Use the `graph` argument instead. + """ + event_writer = EventFileWriter(logdir, max_queue, flush_secs) + super(FileWriter, self).__init__(event_writer, graph, graph_def) + + def get_logdir(self): + """Returns the directory where event file will be written.""" + return self.event_writer.get_logdir() + + def add_event(self, event): + """Adds an event to the event file. + Args: + event: An `Event` protocol buffer. + """ + self.event_writer.add_event(event) + + def flush(self): + """Flushes the event file to disk. + Call this method to make sure that all pending events have been written to + disk. + """ + self.event_writer.flush() + + def close(self): + """Flushes the event file to disk and close the file. + Call this method when you do not need the summary writer anymore. + """ + self.event_writer.close() + + def reopen(self): + """Reopens the EventFileWriter. + Can be called after `close()` to add more events in the same directory. + The events will go into a new events file. + Does nothing if the EventFileWriter was not closed. + """ + self.event_writer.reopen() + + +class SummaryWriter(object): + """Writes `Summary` directly to event files. + The `SummaryWriter` class provides a high-level api to create an event file in a + given directory and add summaries and events to it. The class updates the + file contents asynchronously. This allows a training program to call methods + to add data to the file directly from the training loop, without slowing down + training. + """ + def __init__(self, log_dir=None, comment=''): + """ + Args: + log_dir (string): save location, default is: runs/**CURRENT_DATETIME_HOSTNAME**, which changes after each + run. Use hierarchical folder structure to compare between runs easily. e.g. 'runs/exp1', 'runs/exp2' + comment (string): comment that appends to the default log_dir + """ + if not log_dir: + import socket + from datetime import datetime + current_time = datetime.now().strftime('%b%d_%H-%M-%S') + log_dir = os.path.join('runs', current_time + '_' + socket.gethostname() + comment) + self.file_writer = FileWriter(logdir=log_dir) + v = 1E-12 + buckets = [] + neg_buckets = [] + while v < 1E20: + buckets.append(v) + neg_buckets.append(-v) + v *= 1.1 + self.default_bins = neg_buckets[::-1] + [0] + buckets + self.text_tags = [] + # + self.all_writers = {self.file_writer.get_logdir(): self.file_writer} + self.scalar_dict = {} # {writer_id : [[timestamp, step, value],...],...} + + def __append_to_scalar_dict(self, tag, scalar_value, global_step, + timestamp): + """This adds an entry to the self.scalar_dict datastructure with format + {writer_id : [[timestamp, step, value], ...], ...}. + """ + from .x2num import makenp + if tag not in self.scalar_dict.keys(): + self.scalar_dict[tag] = [] + self.scalar_dict[tag].append([timestamp, global_step, float(makenp(scalar_value))]) + + def add_scalar(self, tag, scalar_value, global_step=None): + """Add scalar data to summary. + + Args: + tag (string): Data identifier + scalar_value (float): Value to save + global_step (int): Global step value to record + """ + self.file_writer.add_summary(scalar(tag, scalar_value), global_step) + self.__append_to_scalar_dict(tag, scalar_value, global_step, time.time()) + + def add_scalars(self, main_tag, tag_scalar_dict, global_step=None): + """Adds many scalar data to summary. + + Args: + tag (string): Data identifier + main_tag (string): The parent name for the tags + tag_scalar_dict (dict): Key-value pair storing the tag and corresponding values + global_step (int): Global step value to record + + Examples:: + + writer.add_scalars('run_14h',{'xsinx':i*np.sin(i/r), + 'xcosx':i*np.cos(i/r), + 'arctanx': numsteps*np.arctan(i/r)}, i) + # This function adds three values to the same scalar plot with the tag + # 'run_14h' in TensorBoard's scalar section. + """ + timestamp = time.time() + fw_logdir = self.file_writer.get_logdir() + for tag, scalar_value in tag_scalar_dict.items(): + fw_tag = fw_logdir + "/" + main_tag + "/" + tag + if fw_tag in self.all_writers.keys(): + fw = self.all_writers[fw_tag] + else: + fw = FileWriter(logdir=fw_tag) + self.all_writers[fw_tag] = fw + fw.add_summary(scalar(main_tag, scalar_value), global_step) + self.__append_to_scalar_dict(fw_tag, scalar_value, global_step, timestamp) + + def export_scalars_to_json(self, path): + """Exports to the given path an ASCII file containing all the scalars written + so far by this instance, with the following format: + {writer_id : [[timestamp, step, value], ...], ...} + """ + with open(path, "w") as f: + json.dump(self.scalar_dict, f) + + def add_histogram(self, tag, values, global_step=None, bins='tensorflow'): + """Add histogram to summary. + + Args: + tag (string): Data identifier + values (numpy.array): Values to build histogram + global_step (int): Global step value to record + bins (string): one of {'tensorflow','auto', 'fd', ...}, this determines how the bins are made. You can find + other options in: https://docs.scipy.org/doc/numpy/reference/generated/numpy.histogram.html + """ + if bins == 'tensorflow': + bins = self.default_bins + self.file_writer.add_summary(histogram(tag, values, bins), global_step) + + def add_image(self, tag, img_tensor, global_step=None): + """Add image data to summary. + + Note that this requires the ``pillow`` package. + + Args: + tag (string): Data identifier + img_tensor (torch.Tensor): Image data + global_step (int): Global step value to record + Shape: + img_tensor: :math:`(3, H, W)`. Use ``torchvision.utils.make_grid()`` to prepare it is a good idea. + """ + self.file_writer.add_summary(image(tag, img_tensor), global_step) + + def add_audio(self, tag, snd_tensor, global_step=None, sample_rate=44100): + """Add audio data to summary. + + Args: + tag (string): Data identifier + snd_tensor (torch.Tensor): Sound data + global_step (int): Global step value to record + sample_rate (int): sample rate in Hz + + Shape: + snd_tensor: :math:`(1, L)`. The values should lie between [-1, 1]. + """ + self.file_writer.add_summary(audio(tag, snd_tensor, sample_rate=sample_rate), global_step) + + def add_text(self, tag, text_string, global_step=None): + """Add text data to summary. + + Args: + tag (string): Data identifier + text_string (string): String to save + global_step (int): Global step value to record + + Examples:: + + writer.add_text('lstm', 'This is an lstm', 0) + writer.add_text('rnn', 'This is an rnn', 10) + """ + self.file_writer.add_summary(text(tag, text_string), global_step) + if tag not in self.text_tags: + self.text_tags.append(tag) + extension_dir = self.file_writer.get_logdir() + '/plugins/tensorboard_text/' + if not os.path.exists(extension_dir): + os.makedirs(extension_dir) + with open(extension_dir + 'tensors.json', 'w') as fp: + json.dump(self.text_tags, fp) + + def add_graph_onnx(self, prototxt): + self.file_writer.add_graph_onnx(gg(prototxt)) + + def add_graph(self, model, input_to_model, verbose=False): + # prohibit second call? + # no, let tensorboard handles it and show its warning message. + """Add graph data to summary. + + Args: + model (torch.nn.Module): model to draw. + input_to_model (torch.autograd.Variable): a variable or a tuple of variables to be fed. + + """ + import torch + from distutils.version import LooseVersion + if LooseVersion(torch.__version__) >= LooseVersion("0.3.1"): + pass + else: + if LooseVersion(torch.__version__) >= LooseVersion("0.3.0"): + print('You are using PyTorch==0.3.0, use add_graph_onnx()') + return + if not hasattr(torch.autograd.Variable, 'grad_fn'): + print('add_graph() only supports PyTorch v0.2.') + return + self.file_writer.add_graph(graph(model, input_to_model, verbose)) + + def add_embedding(self, mat, metadata=None, label_img=None, global_step=None, tag='default'): + """Add embedding projector data to summary. + + Args: + mat (torch.Tensor): A matrix which each row is the feature vector of the data point + metadata (list): A list of labels, each element will be convert to string + label_img (torch.Tensor): Images correspond to each data point + global_step (int): Global step value to record + tag (string): Name for the embedding + Shape: + mat: :math:`(N, D)`, where N is number of data and D is feature dimension + + label_img: :math:`(N, C, H, W)` + + Examples:: + + import keyword + import torch + meta = [] + while len(meta)<100: + meta = meta+keyword.kwlist # get some strings + meta = meta[:100] + + for i, v in enumerate(meta): + meta[i] = v+str(i) + + label_img = torch.rand(100, 3, 10, 32) + for i in range(100): + label_img[i]*=i/100.0 + + writer.add_embedding(torch.randn(100, 5), metadata=meta, label_img=label_img) + writer.add_embedding(torch.randn(100, 5), label_img=label_img) + writer.add_embedding(torch.randn(100, 5), metadata=meta) + """ + if global_step is None: + global_step = 0 + # clear pbtxt? + save_path = os.path.join(self.file_writer.get_logdir(), str(global_step).zfill(5)) + try: + os.makedirs(save_path) + except OSError: + print('warning: Embedding dir exists, did you set global_step for add_embedding()?') + if metadata is not None: + assert mat.size(0) == len(metadata), '#labels should equal with #data points' + make_tsv(metadata, save_path) + if label_img is not None: + assert mat.size(0) == label_img.size(0), '#images should equal with #data points' + make_sprite(label_img, save_path) + assert mat.dim() == 2, 'mat should be 2D, where mat.size(0) is the number of data points' + make_mat(mat.tolist(), save_path) + # new funcion to append to the config file a new embedding + append_pbtxt(metadata, label_img, self.file_writer.get_logdir(), str(global_step).zfill(5), tag) + + def add_pr_curve(self, tag, labels, predictions, global_step=None, num_thresholds=127, weights=None): + """Adds precision recall curve. + + Args: + tag (string): Data identifier + labels (torch.Tensor): Ground truth data. Binary label for each element. + predictions (torch.Tensor): The probability that an element be classified as true. Value should in [0, 1] + global_step (int): Global step value to record + num_thresholds (int): Number of thresholds used to draw the curve. + + """ + from .x2num import makenp + labels = makenp(labels) + predictions = makenp(predictions) + self.file_writer.add_summary(pr_curve(tag, labels, predictions, num_thresholds, weights), global_step) + + def close(self): + if self.file_writer is None: + return # ignore double close + self.file_writer.flush() + self.file_writer.close() + for path, writer in self.all_writers.items(): + writer.flush() + writer.close() + self.file_writer = self.all_writers = None + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() diff --git a/tensorboardX/x2num.py b/tensorboardX/x2num.py new file mode 100644 index 0000000..f2f89d5 --- /dev/null +++ b/tensorboardX/x2num.py @@ -0,0 +1,90 @@ +# DO NOT alter/distruct/free input object ! + +import numpy as np + + +def makenp(x, modality=None): + # if already numpy, return + if isinstance(x, np.ndarray): + if modality == 'IMG' and x.dtype == np.uint8: + return x.astype(np.float32) / 255.0 + return x + if np.isscalar(x): + return np.array([x]) + if 'torch' in str(type(x)): + return pytorch_np(x, modality) + if 'chainer' in str(type(x)): + return chainer_np(x, modality) + if 'mxnet' in str(type(x)): + return mxnet_np(x, modality) + + +def pytorch_np(x, modality): + import torch + if isinstance(x, torch.autograd.Variable): + x = x.data + x = x.cpu().numpy() + if modality == 'IMG': + x = _prepare_image(x) + return x + + +def theano_np(x): + import theano + pass + + +def caffe2_np(x): + pass + + +def mxnet_np(x, modality): + x = x.asnumpy() + if modality == 'IMG': + x = _prepare_image(x) + return x + + +def chainer_np(x, modality): + import chainer + x = chainer.cuda.to_cpu(x.data) + if modality == 'IMG': + x = _prepare_image(x) + return x + + +def make_grid(I, ncols=8): + assert isinstance(I, np.ndarray), 'plugin error, should pass numpy array here' + assert I.ndim == 4 and I.shape[1] == 3 + nimg = I.shape[0] + H = I.shape[2] + W = I.shape[3] + ncols = min(nimg, ncols) + nrows = int(np.ceil(float(nimg) / ncols)) + canvas = np.zeros((3, H * nrows, W * ncols)) + i = 0 + for y in range(nrows): + for x in range(ncols): + if i >= nimg: + break + canvas[:, y * H:(y + 1) * H, x * W:(x + 1) * W] = I[i] + i = i + 1 + return canvas + + +def _prepare_image(I): + assert isinstance(I, np.ndarray), 'plugin error, should pass numpy array here' + assert I.ndim == 2 or I.ndim == 3 or I.ndim == 4 + if I.ndim == 4: # NCHW + if I.shape[1] == 1: # N1HW + I = np.concatenate((I, I, I), 1) # N3HW + assert I.shape[1] == 3 + I = make_grid(I) # 3xHxW + if I.ndim == 3 and I.shape[0] == 1: # 1xHxW + I = np.concatenate((I, I, I), 0) # 3xHxW + if I.ndim == 2: # HxW + I = np.expand_dims(I, 0) # 1xHxW + I = np.concatenate((I, I, I), 0) # 3xHxW + I = I.transpose(1, 2, 0) + + return I diff --git a/test.py b/test.py index 3e60ae7..f081d52 100644 --- a/test.py +++ b/test.py @@ -15,8 +15,9 @@ from PIL import Image import lib.TestTransforms as transforms from models.ColorVidNet import ColorVidNet from models.FrameColor import frame_colorization -from models.NonlocalNet import VGG19_feature_color_torchversion, WarpNet -from utils.util import batch_lab2rgb_transpose_mc, folder2vid, mkdir_if_not, save_frames, tensor_lab2rgb, uncenter_l +from models.NonlocalNet import VGG19_pytorch, WarpNet +from utils.util import (batch_lab2rgb_transpose_mc, folder2vid, mkdir_if_not, + save_frames, tensor_lab2rgb, uncenter_l) from utils.util_distortion import CenterPad, Normalize, RGB2Lab, ToTensor os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" @@ -36,7 +37,7 @@ def colorize_video(opt, input_path, reference_file, output_path, nonlocal_net, c print("processing the folder:", input_path) path, dirs, filenames = os.walk(input_path).__next__() file_count = len(filenames) - filenames.sort() + filenames.sort(key=lambda f: int("".join(filter(str.isdigit, f) or -1))) # NOTE: resize frames to 216*384 transform = transforms.Compose( @@ -129,6 +130,7 @@ def colorize_video(opt, input_path, reference_file, output_path, nonlocal_net, c # save the frames save_frames(IA_predict_rgb, output_path, index) + # output video video_name = "video.avi" folder2vid(image_folder=output_path, output_dir=output_path, filename=video_name) @@ -145,7 +147,7 @@ if __name__ == "__main__": parser.add_argument("--clip_path", type=str, default="./sample_videos/clips/v32", help="path of input clips") parser.add_argument("--ref_path", type=str, default="./sample_videos/ref/v32", help="path of refernce images") parser.add_argument("--output_path", type=str, default="./sample_videos/output", help="path of output clips") - opt = parser.parse_args(args=[]) + opt = parser.parse_args() opt.gpu_ids = [int(x) for x in opt.gpu_ids.split(",")] cudnn.benchmark = True print("running on GPU", opt.gpu_ids) @@ -156,7 +158,7 @@ if __name__ == "__main__": nonlocal_net = WarpNet(1) colornet = ColorVidNet(7) - vggnet = VGG19_feature_color_torchversion() + vggnet = VGG19_pytorch() vggnet.load_state_dict(torch.load("data/vgg19_conv.pth")) for param in vggnet.parameters(): param.requires_grad = False diff --git a/train.py b/train.py new file mode 100644 index 0000000..349fbc0 --- /dev/null +++ b/train.py @@ -0,0 +1,738 @@ +from __future__ import print_function + +import argparse +import math +import os +import queue +import time + +import cv2 +import torch +import torch.backends.cudnn as cudnn +import torch.nn as nn +import torch.optim as optim +import torchvision.transforms as transform_lib +from torch.utils.data import ConcatDataset, DataLoader, WeightedRandomSampler +from torchvision.transforms import CenterCrop + +import lib.TrainTransforms as transforms +from lib.videoloader import VideosDataset +from lib.videoloader_imagenet import VideosDataset_ImageNet +from models.ColorVidNet import ColorVidNet +from models.ContextualLoss import ContextualLoss, ContextualLoss_forward +from models.FrameColor import frame_colorization +from models.GAN_models import Discriminator_x64 +from models.NonlocalNet import (NonlocalWeightedAverage, VGG19_pytorch, + WarpNet, WeightedAverage, + WeightedAverage_color) +from tensorboardX import SummaryWriter +from utils.util import (batch_lab2rgb_transpose_mc, feature_normalize, l1_loss, + mkdir_if_not, mse_loss, parse, tensor_lab2rgb, + uncenter_l, weighted_l1_loss, weighted_mse_loss) +from utils.util_distortion import (CenterPad_threshold, Normalize, RGB2Lab, + ToTensor) +from utils.util_tensorboard import TBImageRecorder, value_logger +from utils.warping import WarpingLayer + +cv2.setNumThreads(0) + +parser = argparse.ArgumentParser() +parser.add_argument("--data_root", default="/home/zhanbo/remote/video/video_pair3/", type=str) +parser.add_argument("--data_root_imagenet", default="/home/zhanbo/remote/imagenet_ref_pair/", type=str) +parser.add_argument("--gpu_ids", type=str, default="0,1,2,3", help="separate by comma") +parser.add_argument("--workers", type=int, default=8) +parser.add_argument("--batch_size", type=int, default=16) +parser.add_argument("--image_size", type=int, default=[216, 384]) +parser.add_argument("--ic", type=int, default=7) +parser.add_argument("--epoch", type=int, default=40) + +parser.add_argument("--resume_epoch", type=int, default=0) +parser.add_argument("--resume", type=bool, default=False) +parser.add_argument("--load_pretrained_model", type=bool, default=True) + +parser.add_argument("--lr", type=float, default=0.00001) +parser.add_argument("--beta1", type=float, default=0.5) +parser.add_argument("--lr_step", type=int, default=100) +parser.add_argument("--lr_gamma", type=float, default=0.1) + + +parser.add_argument("--checkpoint_dir", type=str, default="checkpoints/default") +parser.add_argument("--tb_log_step", type=int, default=50) +parser.add_argument("--print_step", type=int, default=2) + +parser.add_argument("--real_reference_probability", type=float, default=0.7) +parser.add_argument("--nonzero_placeholder_probability", type=float, default=0.0) +parser.add_argument("--with_bad", type=bool, default=True) +parser.add_argument("--with_mid", type=bool, default=True) + +parser.add_argument("--domain_invariant", type=bool, default=False) +parser.add_argument("--weigth_l1", type=float, default=2.0) +parser.add_argument("--weight_contextual", type=float, default="0.2") +parser.add_argument("--weight_perceptual", type=float, default="0.001") +parser.add_argument("--weight_smoothness", type=float, default="5.0") +parser.add_argument("--weight_gan", type=float, default="0.2") +parser.add_argument("--weight_nonlocal_smoothness", type=float, default="0.0") +parser.add_argument("--weight_nonlocal_consistent", type=float, default="0.0") +parser.add_argument("--weight_consistent", type=float, default="0.02") +parser.add_argument("--luminance_noise", type=float, default="2.0") +parser.add_argument("--permute_data", type=bool, default=True) +parser.add_argument("--contextual_loss_direction", type=str, default="forward", help="forward or backward matching") + + +def image_logger_fn( + I_last_lab, + I_current_lab, + I_reference_lab, + I_last_lab_predict, + I_current_lab_predict, + I_last_nonlocal_lab, + I_current_nonlocal_lab, +): + I_last_image = batch_lab2rgb_transpose_mc(I_last_lab[0:32, 0:1, :, :], I_last_lab[0:32, 1:3, :, :]) + I_current_image = batch_lab2rgb_transpose_mc(I_current_lab[0:32, 0:1, :, :], I_current_lab[0:32, 1:3, :, :]) + I_reference_image = batch_lab2rgb_transpose_mc(I_reference_lab[0:32, 0:1, :, :], I_reference_lab[0:32, 1:3, :, :]) + I_last_image_predict = batch_lab2rgb_transpose_mc( + I_last_lab_predict[0:32, 0:1, :, :], I_last_lab_predict[0:32, 1:3, :, :] + ) + I_current_image_predict = batch_lab2rgb_transpose_mc( + I_current_lab_predict[0:32, 0:1, :, :], I_current_lab_predict[0:32, 1:3, :, :] + ) + I_last_nonlocal_image = batch_lab2rgb_transpose_mc( + I_last_nonlocal_lab[0:32, 0:1, :, :], I_last_nonlocal_lab[0:32, 1:3, :, :] + ) + I_current_nonlocal_image = batch_lab2rgb_transpose_mc( + I_current_nonlocal_lab[0:32, 0:1, :, :], I_current_nonlocal_lab[0:32, 1:3, :, :] + ) + + img_info = {} + img_info["1_I_last"] = I_last_image + img_info["2_I_current"] = I_current_image + img_info["3_I_reference"] = I_reference_image + img_info["4_I_last_predict"] = I_last_image_predict + img_info["5_I_curren_predict"] = I_current_image_predict + img_info["6_I_last_nonlocal"] = I_last_nonlocal_image + img_info["7_I_current_nonlocal"] = I_current_nonlocal_image + + return img_info + + +def training_logger(): + try: + if total_iter % opt.print_step == 0: + print("processing time:", elapsed) + print( + "Epoch %d, Step[%d/%d], lr: %f, total_loss: %.2f" + % ( + epoch, + ((iter + 1) % iter_num_per_epoch), + iter_num_per_epoch, + step_optim_scheduler_g.get_last_lr()[0], + total_loss.item(), + ) + ) + + value_logger( + tb_writer, + total_iter, + loss_info={ + "l1_loss": l1_loss.item(), + "feat_loss": feat_loss.item(), + "contextual_loss_total": contextual_loss_total.item(), + "smoothness_loss": smoothness_loss.item(), + "nonlocal_smoothness_loss": nonlocal_smoothness_loss.item(), + "nonlocal_consistent_loss": nonlocal_consistent_loss.item(), + "consistent_loss": consistent_loss.item(), + "generator_loss": generator_loss.item(), + "discriminator_loss": discriminator_loss.item(), + "total_loss": total_loss.item(), + }, + ) + + if total_iter % opt.tb_log_step == 0: + I_last_nonlocal_lab = torch.cat( + (I_last_lab[:, 0:1, :, :], I_last_nonlocal_lab_predict[:, 1:3, :, :]), dim=1 + ) + I_current_nonlocal_lab = torch.cat( + (I_current_lab[:, 0:1, :, :], I_current_nonlocal_lab_predict[:, 1:3, :, :]), dim=1 + ) + I_last_lab_predict = torch.cat((I_last_l, I_last_ab_predict), dim=1) + data_queue.put( + ( + ( + I_last_lab.cpu(), + I_current_lab.cpu(), + I_reference_lab.cpu(), + I_last_lab_predict.cpu(), + I_current_lab_predict.cpu(), + I_last_nonlocal_lab.cpu(), + I_current_nonlocal_lab, + ), + total_iter, + ) + ) + + if total_iter % 2000 == 0: + if len(opt.gpu_ids) > 1: + torch.save( + nonlocal_net.module.state_dict(), + os.path.join(opt.checkpoint_dir, "nonlocal_net_iter_%d.pth") % total_iter, + ) + torch.save( + colornet.module.state_dict(), + os.path.join(opt.checkpoint_dir, "colornet_iter_%d.pth") % total_iter, + ) + torch.save( + discriminator.module.state_dict(), + os.path.join(opt.checkpoint_dir, "discriminator_iter_%d.pth") % total_iter, + ) + else: + torch.save( + nonlocal_net.state_dict(), + os.path.join(opt.checkpoint_dir, "nonlocal_net_iter_%d.pth") % total_iter, + ) + torch.save(colornet.state_dict(), os.path.join(opt.checkpoint_dir, "colornet_iter_%d.pth") % total_iter) + torch.save( + discriminator.state_dict(), + os.path.join(opt.checkpoint_dir, "discriminator_iter_%d.pth") % total_iter, + ) + + # save the state for resume + if total_iter % 2000 == 0: + print("saving the checkpoint") + if len(opt.gpu_ids) > 1: + state = { + "total_iter": total_iter + 1, + "epoch": epoch, + "colornet_state": colornet.module.state_dict(), + "nonlocal_net_state": nonlocal_net.module.state_dict(), + "discriminator_state": discriminator.module.state_dict(), + "optimizer_g": optimizer_g.state_dict(), + "optimizer_d": optimizer_d.state_dict(), + "optimizer_schedule_g": step_optim_scheduler_g.state_dict(), + "optimizer_schedule_d": step_optim_scheduler_g.state_dict(), + } + else: + state = { + "total_iter": total_iter + 1, + "epoch": epoch, + "colornet_state": colornet.state_dict(), + "nonlocal_net_state": nonlocal_net.state_dict(), + "discriminator_state": discriminator.state_dict(), + "optimizer_g": optimizer_g.state_dict(), + "optimizer_d": optimizer_d.state_dict(), + "optimizer_schedule_g": step_optim_scheduler_g.state_dict(), + "optimizer_schedule_d": step_optim_scheduler_d.state_dict(), + } + torch.save(state, os.path.join(opt.checkpoint_dir, "learning_checkpoint.pth")) + except Exception as e: + print("Exception during output") + print(e) + + +def gpu_setup(): + os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" + cudnn.benchmark = True + opt.gpu_ids = [int(x) for x in opt.gpu_ids.split(",")] + torch.cuda.set_device(opt.gpu_ids[0]) + device = torch.device("cuda") + print("running on GPU", opt.gpu_ids) + return device + + +def load_data(): + print("initializing dataloader") + transforms_video = [ + CenterCrop(opt.image_size), + RGB2Lab(), + ToTensor(), + Normalize(), + ] + transforms_imagenet = [CenterPad_threshold(opt.image_size), RGB2Lab(), ToTensor(), Normalize()] + extra_reference_transform = [ + transform_lib.RandomHorizontalFlip(0.5), + transform_lib.RandomResizedCrop(480, (0.98, 1.0), ratio=(0.8, 1.2)), + ] + train_dataset_video = VideosDataset( + data_root=opt.data_root, + epoch=opt.epoch, + image_size=opt.image_size, + image_transform=transforms.Compose(transforms_video), + real_reference_probability=opt.real_reference_probability, + nonzero_placeholder_probability=opt.nonzero_placeholder_probability, + ) + train_dataset_imagenet = VideosDataset_ImageNet( + data_root=opt.data_root_imagenet, + image_size=opt.image_size, + epoch=opt.epoch, + with_bad=opt.with_bad, + with_mid=opt.with_mid, + transforms_imagenet=transforms_imagenet, + distortion_level=4, + brightnessjitter=5, + nonzero_placeholder_probability=opt.nonzero_placeholder_probability, + extra_reference_transform=extra_reference_transform, + real_reference_probability=opt.real_reference_probability, + ) + + video_training_length = len(train_dataset_video) + imagenet_training_length = len(train_dataset_imagenet) + dataset_training_length = train_dataset_video.real_len + train_dataset_imagenet.real_len + dataset_combined = ConcatDataset([train_dataset_video, train_dataset_imagenet]) + sampler = WeightedRandomSampler( + [1] * video_training_length + [1] * imagenet_training_length, dataset_training_length * opt.epoch + ) + data_loader = DataLoader( + dataset_combined, + batch_size=opt.batch_size, + shuffle=False, + num_workers=opt.workers, + pin_memory=True, + drop_last=True, + sampler=sampler, + ) + return dataset_training_length, train_dataset_video, train_dataset_imagenet, data_loader + + +def define_loss(): + print("defining loss") + ab_criterion = nn.SmoothL1Loss().to(device) + nonlocal_criterion = nn.SmoothL1Loss().to(device) + feat_l2_criterion = nn.MSELoss().to(device) + feat_l1_criterion = nn.SmoothL1Loss().to(device) + contextual_loss = ContextualLoss().to(device) + contextual_forward_loss = ContextualLoss_forward().to(device) + BCE_stable = nn.BCEWithLogitsLoss().to(device) + return contextual_loss, contextual_forward_loss + + +def define_optimizer(): + print("defining optimizer") + optimizer_g = optim.Adam( + [{"params": nonlocal_net.parameters(), "lr": 1e-5}, {"params": colornet.parameters(), "lr": 2e-4}], + betas=(0.5, 0.999), + eps=1e-5, + amsgrad=True, + ) + optimizer_d = optim.Adam( + filter(lambda p: p.requires_grad, discriminator.parameters()), lr=2 * 1e-4, betas=(0.5, 0.999) + ) + return optimizer_g, optimizer_d + + +def resume_model(): + print("resuming the learning") + checkpoint = torch.load(os.path.join(opt.checkpoint_dir, "learning_checkpoint.pth")) + total_iter = checkpoint["total_iter"] + epoch = checkpoint["epoch"] + colornet.load_state_dict(checkpoint["colornet_state"]) + nonlocal_net.load_state_dict(checkpoint["nonlocal_net_state"]) + discriminator.load_state_dict(checkpoint["discriminator_state"]) + optimizer_g.load_state_dict(checkpoint["optimizer_g"]) + optimizer_d.load_state_dict(checkpoint["optimizer_d"]) + step_optim_scheduler_g.load_state_dict(checkpoint["optimizer_schedule_g"]) + step_optim_scheduler_d.load_state_dict(checkpoint["optimizer_schedule_d"]) + + +def to_device( + colornet, + nonlocal_net, + discriminator, + vggnet, + contextual_loss, + contextual_forward_loss, + weighted_layer_color, + nonlocal_weighted_layer, + warping_layer, + instancenorm, +): + print("moving models to device") + colornet = torch.nn.DataParallel(colornet.to(device), device_ids=opt.gpu_ids) + nonlocal_net = torch.nn.DataParallel(nonlocal_net.to(device), device_ids=opt.gpu_ids) + discriminator = torch.nn.DataParallel(discriminator.to(device), device_ids=opt.gpu_ids) + vggnet = torch.nn.DataParallel(vggnet.to(device), device_ids=opt.gpu_ids) + contextual_loss = torch.nn.DataParallel(contextual_loss.to(device), device_ids=opt.gpu_ids) + contextual_forward_loss = torch.nn.DataParallel(contextual_forward_loss.to(device), device_ids=opt.gpu_ids) + weighted_layer_color = torch.nn.DataParallel(weighted_layer_color.to(device), device_ids=opt.gpu_ids) + nonlocal_weighted_layer = torch.nn.DataParallel(nonlocal_weighted_layer.to(device), device_ids=opt.gpu_ids) + warping_layer = torch.nn.DataParallel(warping_layer.to(device), device_ids=opt.gpu_ids) + instancenorm = torch.nn.DataParallel(instancenorm.to(device), device_ids=opt.gpu_ids) + + return ( + vggnet, + nonlocal_net, + colornet, + discriminator, + instancenorm, + contextual_loss, + contextual_forward_loss, + weighted_layer_color, + nonlocal_weighted_layer, + warping_layer, + ) + + +def loss_init(): + print("initializing losses") + zero_loss = torch.Tensor([0]).to(device) + ( + feat_loss, + contextual_loss_total, + smoothness_loss, + nonlocal_smoothness_loss, + consistent_loss, + nonlocal_consistent_loss, + generator_loss, + discriminator_loss, + ) = (zero_loss, zero_loss, zero_loss, zero_loss, zero_loss, zero_loss, zero_loss, zero_loss) + + return ( + feat_loss, + contextual_loss_total, + smoothness_loss, + nonlocal_smoothness_loss, + consistent_loss, + nonlocal_consistent_loss, + generator_loss, + discriminator_loss, + ) + + +def video_colorization(): + # colorization for the last frame + I_last_ab_predict, I_last_nonlocal_lab_predict, features_last_gray = frame_colorization( + I_last_lab, + I_reference_lab, + placeholder_lab, + features_B, + vggnet, + nonlocal_net, + colornet, + feature_noise=0, + luminance_noise=opt.luminance_noise, + ) + I_last_lab_predict = torch.cat((I_last_l, I_last_ab_predict), dim=1) + + # colorization for the current frame + I_current_ab_predict, I_current_nonlocal_lab_predict, features_current_gray = frame_colorization( + I_current_lab, + I_reference_lab, + I_last_lab_predict, + features_B, + vggnet, + nonlocal_net, + colornet, + feature_noise=0, + luminance_noise=opt.luminance_noise, + ) + I_current_lab_predict = torch.cat((I_last_l, I_current_ab_predict), dim=1) + return I_current_ab_predict, I_last_ab_predict, I_current_nonlocal_lab_predict, I_last_nonlocal_lab_predict + + +if __name__ == "__main__": + torch.multiprocessing.set_start_method("spawn", force=True) + opt = parse(parser) + opt.data_root = opt.data_root.split(",")[0] + opt.data_root_imagenet = opt.data_root_imagenet.split(",")[0] + mkdir_if_not(opt.checkpoint_dir) + mkdir_if_not("./runs/") + + device = gpu_setup() + dataset_training_length, train_dataset_video, train_dataset_imagenet, data_loader = load_data() + tb_writer = SummaryWriter() + data_queue = queue.Queue() + tb_image_reorder = TBImageRecorder(tb_writer, image_logger_fn, data_queue) + tb_image_reorder.start() + + # define network + nonlocal_net = WarpNet(opt.batch_size) + colornet = ColorVidNet(opt.ic) + discriminator = Discriminator_x64(ndf=64) + colornet.to(device) + nonlocal_net.to(device) + discriminator.to(device) + + weighted_layer = WeightedAverage() + weighted_layer_color = WeightedAverage_color() + nonlocal_weighted_layer = NonlocalWeightedAverage() + instancenorm = nn.InstanceNorm2d(512, affine=False) + warping_layer = WarpingLayer("gpu") + + vggnet = VGG19_pytorch() + vggnet.load_state_dict(torch.load("data/vgg19_conv.pth")) + vggnet.eval() + for param in vggnet.parameters(): + param.requires_grad = False + + # load pre-trained model + if opt.load_pretrained_model: + nonlocal_pretain_path = os.path.join("checkpoints/video_moredata_l1/", "nonlocal_net_iter_76000.pth") + nonlocal_net.load_state_dict(torch.load(nonlocal_pretain_path)) + color_test_path = "checkpoints/video_moredata_l1/" + "colornet_iter_76000.pth" + colornet.load_state_dict(torch.load(color_test_path)) + + # define loss function + contextual_loss, contextual_forward_loss = define_loss() + + # define optimizer + optimizer_g, optimizer_d = define_optimizer() + step_optim_scheduler_g = optim.lr_scheduler.StepLR(optimizer_g, step_size=opt.lr_step, gamma=opt.lr_gamma) + step_optim_scheduler_d = optim.lr_scheduler.StepLR(optimizer_d, step_size=opt.lr_step, gamma=opt.lr_gamma) + + # define others + downsampling_by2 = nn.AvgPool2d(kernel_size=2).to(device) + downsampling_by4 = nn.AvgPool2d(kernel_size=4).to(device) + + # dataset info + iter_num_per_epoch = dataset_training_length // opt.batch_size + total_iter = opt.resume_epoch * iter_num_per_epoch + print( + "train_dataset info, real_len: %d, epoch_len: %d, iter_num_per_epoch: %d" + % (dataset_training_length, len(train_dataset_video) + len(train_dataset_imagenet), iter_num_per_epoch) + ) + + if opt.resume: + resume_model() + + # move to GPU processing + device = opt.gpu_ids[0] + ( + vggnet, + nonlocal_net, + colornet, + discriminator, + instancenorm, + contextual_loss, + contextual_forward_loss, + weighted_layer_color, + nonlocal_weighted_layer, + warping_layer, + ) = to_device( + colornet, + nonlocal_net, + discriminator, + vggnet, + contextual_loss, + contextual_forward_loss, + weighted_layer_color, + nonlocal_weighted_layer, + warping_layer, + instancenorm, + ) + + ( + feat_loss, + contextual_loss_total, + smoothness_loss, + nonlocal_smoothness_loss, + consistent_loss, + nonlocal_consistent_loss, + generator_loss, + discriminator_loss, + ) = loss_init() + + # %% Training + print("start training") + for iter, data in enumerate(data_loader): + start_time = time.time() + total_iter += 1 + epoch = math.ceil(total_iter / iter_num_per_epoch) + + ###### LOADING DATA SAMPLE ###### + ( + I_last_lab, + I_current_lab, + I_reference_lab, + flow_forward, + flow_backward, + mask, + placeholder_lab, + self_ref_flag, + ) = data + I_last_lab = I_last_lab.to(device) + I_current_lab = I_current_lab.to(device) + I_reference_lab = I_reference_lab.to(device) + flow_forward = flow_forward.to(device) + flow_backward = flow_backward.to(device) + mask = mask.to(device) + placeholder_lab = placeholder_lab.to(device) + self_ref_flag = self_ref_flag.to(device) + + I_last_l = I_last_lab[:, 0:1, :, :] + I_last_ab = I_last_lab[:, 1:3, :, :] + I_current_l = I_current_lab[:, 0:1, :, :] + I_current_ab = I_current_lab[:, 1:3, :, :] + I_reference_l = I_reference_lab[:, 0:1, :, :] + I_reference_ab = I_reference_lab[:, 1:3, :, :] + I_reference_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_reference_l), I_reference_ab), dim=1)) + features_B = vggnet(I_reference_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True) + + ###### COLORIZATION ###### + ( + I_current_ab_predict, + I_last_ab_predict, + I_current_nonlocal_lab_predict, + I_last_nonlocal_lab_predict, + ) = video_colorization() + + ###### UPDATE DISCRIMINATOR ###### + if opt.weight_gan > 0: + optimizer_g.zero_grad() + optimizer_d.zero_grad() + + fake_data_lab = torch.cat( + (uncenter_l(I_current_l), I_current_ab_predict, uncenter_l(I_last_l), I_last_ab_predict), dim=1 + ) + real_data_lab = torch.cat((uncenter_l(I_current_l), I_current_ab, uncenter_l(I_last_l), I_last_ab), dim=1) + + if opt.permute_data: + batch_index = torch.arange(-1, opt.batch_size - 1, dtype=torch.long) + real_data_lab = real_data_lab[batch_index, ...] + + y_pred_fake, feature_pred_fake = discriminator(fake_data_lab.detach()) + y_pred_real, feature_pred_real = discriminator(real_data_lab.detach()) + + y = torch.ones_like(y_pred_real) + y2 = torch.zeros_like(y_pred_real) + discriminator_loss = ( + torch.mean((y_pred_real - torch.mean(y_pred_fake) - y) ** 2) + + torch.mean((y_pred_fake - torch.mean(y_pred_real) + y) ** 2) + ) / 2 + discriminator_loss.backward() + optimizer_d.step() + + ###### UPDATE GENERATOR ###### + optimizer_g.zero_grad() + optimizer_d.zero_grad() + + # extract vgg features for both output and original image + I_predict_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_current_l), I_current_ab_predict), dim=1)) + predict_relu1_1, predict_relu2_1, predict_relu3_1, predict_relu4_1, predict_relu5_1 = vggnet( + I_predict_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True + ) + + I_current_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_current_l), I_current_ab), dim=1)) + A_relu1_1, A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1 = vggnet( + I_current_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True + ) + B_relu1_1, B_relu2_1, B_relu3_1, B_relu4_1, B_relu5_1 = features_B + + ###### LOSS COMPUTE ###### + # l1 loss + if opt.weigth_l1 > 0: + sample_weights = (self_ref_flag[:, 1:3, :, :]) / (sum(self_ref_flag[:, 0, 0, 0]) + 1e-5) + l1_loss = weighted_l1_loss(I_current_ab_predict, I_current_ab, sample_weights) * opt.weigth_l1 + + # generator loss + if opt.weight_gan > 0: + y_pred_fake, feature_pred_fake = discriminator(fake_data_lab) + y_pred_real, feature_pred_real = discriminator(real_data_lab) + generator_loss = ( + ( + torch.mean((y_pred_real - torch.mean(y_pred_fake) + y) ** 2) + + torch.mean((y_pred_fake - torch.mean(y_pred_real) - y) ** 2) + ) + / 2 + * opt.weight_gan + ) + + # feature loss + if opt.domain_invariant: + feat_loss = ( + mse_loss(instancenorm(predict_relu5_1), instancenorm(A_relu5_1.detach())) + * opt.weight_perceptual + * 1e5 + * 0.2 + ) + else: + feat_loss = mse_loss(predict_relu5_1, A_relu5_1.detach()) * opt.weight_perceptual + + # contextual loss + if opt.contextual_loss_direction == "backward": + contextual_style5_1 = torch.mean(contextual_loss(predict_relu5_1, B_relu5_1.detach())) * 8 + contextual_style4_1 = torch.mean(contextual_loss(predict_relu4_1, B_relu4_1.detach())) * 4 + contextual_style3_1 = ( + torch.mean(contextual_loss(downsampling_by2(predict_relu3_1), downsampling_by2(B_relu3_1.detach()))) * 2 + ) + else: + contextual_style5_1 = torch.mean(contextual_forward_loss(predict_relu5_1, B_relu5_1.detach())) * 8 + contextual_style4_1 = torch.mean(contextual_forward_loss(predict_relu4_1, B_relu4_1.detach())) * 4 + contextual_style3_1 = ( + torch.mean( + contextual_forward_loss(downsampling_by2(predict_relu3_1), downsampling_by2(B_relu3_1.detach())) + ) + * 2 + ) + if opt.weight_contextual > 0: + contextual_loss_total = ( + contextual_style5_1 + contextual_style4_1 + contextual_style3_1 + ) * opt.weight_contextual + + # smoothness loss + if opt.weight_smoothness > 0: + scale_factor = 1 + I_current_lab_predict = torch.cat((I_current_l, I_current_ab_predict), dim=1) + IA_ab_weighed = weighted_layer_color( + I_current_lab, I_current_lab_predict, patch_size=3, alpha=10, scale_factor=scale_factor + ) + smoothness_loss = ( + mse_loss(nn.functional.interpolate(I_current_ab_predict, scale_factor=scale_factor), IA_ab_weighed) + * opt.weight_smoothness + ) + + if opt.weight_nonlocal_smoothness > 0: + scale_factor = 0.25 + alpha_nonlocal_smoothness = 0.5 + nonlocal_smooth_feature = feature_normalize(A_relu2_1) + I_current_lab_predict = torch.cat((I_current_l, I_current_ab_predict), dim=1) + I_current_ab_weighted_nonlocal = nonlocal_weighted_layer( + I_current_lab_predict, + nonlocal_smooth_feature.detach(), + patch_size=3, + alpha=alpha_nonlocal_smoothness, + scale_factor=scale_factor, + ) + nonlocal_smoothness_loss = ( + mse_loss( + nn.functional.interpolate(I_current_ab_predict, scale_factor=scale_factor), + I_current_ab_weighted_nonlocal, + ) + * opt.weight_nonlocal_smoothness + ) + + if opt.weight_consistent: + I_current_lab_predict_warp = warping_layer(I_current_lab_predict, flow_forward) + I_current_ab_predict_warp = I_current_lab_predict_warp[:, 1:3, :, :] + consistent_loss = ( + weighted_mse_loss(I_current_ab_predict_warp, I_last_ab_predict, mask) * opt.weight_consistent + ) + + if opt.weight_nonlocal_consistent: + I_current_nonlocal_lab_predict_warp = warping_layer(I_current_nonlocal_lab_predict, flow_forward) + nonlocal_consistent_loss = ( + weighted_mse_loss( + I_current_nonlocal_lab_predict_warp[:, 1:3, :, :], I_last_nonlocal_lab_predict[:, 1:3, :, :], mask + ) + * opt.weight_nonlocal_consistent + ) + + # total loss + total_loss = ( + l1_loss + + feat_loss + + contextual_loss_total + + smoothness_loss + + nonlocal_smoothness_loss + + consistent_loss + + nonlocal_consistent_loss + + generator_loss + ) + total_loss.backward() + optimizer_g.step() + end_time = time.time() + elapsed = end_time - start_time + training_logger() + + step_optim_scheduler_g.step() + step_optim_scheduler_d.step() + + data_queue.put((None, None)) diff --git a/utils/flowlib.py b/utils/flowlib.py new file mode 100644 index 0000000..5dae596 --- /dev/null +++ b/utils/flowlib.py @@ -0,0 +1,523 @@ +#!/usr/bin/python +""" +# ============================== +# flowlib.py +# library for optical flow processing +# Author: Ruoteng Li +# Date: 6th Aug 2016 +# ============================== +""" +import matplotlib.colors as cl +import matplotlib.pyplot as plt +import numpy as np +import png +from PIL import Image + +UNKNOWN_FLOW_THRESH = 1e7 +SMALLFLOW = 0.0 +LARGEFLOW = 1e8 + +""" +============= +Flow Section +============= +""" + + +def show_flow(filename): + """ + visualize optical flow map using matplotlib + :param filename: optical flow file + :return: None + """ + flow = read_flow(filename) + img = flow_to_image(flow) + plt.imshow(img) + plt.show() + + +def visualize_flow(flow, mode="Y"): + """ + this function visualize the input flow + :param flow: input flow in array + :param mode: choose which color mode to visualize the flow (Y: Ccbcr, RGB: RGB color) + :return: None + """ + if mode == "Y": + # Ccbcr color wheel + img = flow_to_image(flow) + plt.imshow(img) + plt.show() + elif mode == "RGB": + (h, w) = flow.shape[0:2] + du = flow[:, :, 0] + dv = flow[:, :, 1] + valid = flow[:, :, 2] + max_flow = max(np.max(du), np.max(dv)) + img = np.zeros((h, w, 3), dtype=np.float64) + # angle layer + img[:, :, 0] = np.arctan2(dv, du) / (2 * np.pi) + # magnitude layer, normalized to 1 + img[:, :, 1] = np.sqrt(du * du + dv * dv) * 8 / max_flow + # phase layer + img[:, :, 2] = 8 - img[:, :, 1] + # clip to [0,1] + small_idx = img[:, :, 0:3] < 0 + large_idx = img[:, :, 0:3] > 1 + img[small_idx] = 0 + img[large_idx] = 1 + # convert to rgb + img = cl.hsv_to_rgb(img) + # remove invalid point + img[:, :, 0] = img[:, :, 0] * valid + img[:, :, 1] = img[:, :, 1] * valid + img[:, :, 2] = img[:, :, 2] * valid + # show + plt.imshow(img) + plt.show() + + return None + + +def read_flow(filename): + """ + read optical flow from Middlebury .flo file + :param filename: name of the flow file + :return: optical flow data in matrix + """ + f = open(filename, "rb") + try: + magic = np.fromfile(f, np.float32, count=1)[0] # For Python3.x + except: + magic = np.fromfile(f, np.float32, count=1) # For Python2.x + data2d = None + + if 202021.25 != magic: + print("Magic number incorrect. Invalid .flo file") + else: + w = np.fromfile(f, np.int32, count=1)[0] + h = np.fromfile(f, np.int32, count=1)[0] + # print("Reading %d x %d flo file" % (h, w)) + data2d = np.fromfile(f, np.float32, count=2 * w * h) + # reshape data into 3D array (columns, rows, channels) + data2d = np.resize(data2d, (h, w, 2)) + f.close() + return data2d + + +def read_flow_png(flow_file): + """ + Read optical flow from KITTI .png file + :param flow_file: name of the flow file + :return: optical flow data in matrix + """ + flow_object = png.Reader(filename=flow_file) + flow_direct = flow_object.asDirect() + flow_data = list(flow_direct[2]) + (w, h) = flow_direct[3]["size"] + flow = np.zeros((h, w, 3), dtype=np.float64) + for i in range(len(flow_data)): + flow[i, :, 0] = flow_data[i][0::3] + flow[i, :, 1] = flow_data[i][1::3] + flow[i, :, 2] = flow_data[i][2::3] + + invalid_idx = flow[:, :, 2] == 0 + flow[:, :, 0:2] = (flow[:, :, 0:2] - 2 ** 15) / 64.0 + flow[invalid_idx, 0] = 0 + flow[invalid_idx, 1] = 0 + return flow + + +def write_flow(flow, filename): + """ + write optical flow in Middlebury .flo format + :param flow: optical flow map + :param filename: optical flow file path to be saved + :return: None + """ + with open(filename, "wb") as f: + magic = np.array([202021.25], dtype=np.float32) + (height, width) = flow.shape[0:2] + w = np.array([width], dtype=np.int32) + h = np.array([height], dtype=np.int32) + magic.tofile(f) + w.tofile(f) + h.tofile(f) + flow.tofile(f) + + +def segment_flow(flow): + h = flow.shape[0] + w = flow.shape[1] + u = flow[:, :, 0] + v = flow[:, :, 1] + + idx = (abs(u) > LARGEFLOW) | (abs(v) > LARGEFLOW) + idx2 = abs(u) == SMALLFLOW + class0 = (v == 0) & (u == 0) + u[idx2] = 0.00001 + tan_value = v / u + + class1 = (tan_value < 1) & (tan_value >= 0) & (u > 0) & (v >= 0) + class2 = (tan_value >= 1) & (u >= 0) & (v >= 0) + class3 = (tan_value < -1) & (u <= 0) & (v >= 0) + class4 = (tan_value < 0) & (tan_value >= -1) & (u < 0) & (v >= 0) + class8 = (tan_value >= -1) & (tan_value < 0) & (u > 0) & (v <= 0) + class7 = (tan_value < -1) & (u >= 0) & (v <= 0) + class6 = (tan_value >= 1) & (u <= 0) & (v <= 0) + class5 = (tan_value >= 0) & (tan_value < 1) & (u < 0) & (v <= 0) + + seg = np.zeros((h, w)) + + seg[class1] = 1 + seg[class2] = 2 + seg[class3] = 3 + seg[class4] = 4 + seg[class5] = 5 + seg[class6] = 6 + seg[class7] = 7 + seg[class8] = 8 + seg[class0] = 0 + seg[idx] = 0 + + return seg + + +def flow_error(tu, tv, u, v): + """ + Calculate average end point error + :param tu: ground-truth horizontal flow map + :param tv: ground-truth vertical flow map + :param u: estimated horizontal flow map + :param v: estimated vertical flow map + :return: End point error of the estimated flow + """ + smallflow = 0.0 + """ + stu = tu[bord+1:end-bord,bord+1:end-bord] + stv = tv[bord+1:end-bord,bord+1:end-bord] + su = u[bord+1:end-bord,bord+1:end-bord] + sv = v[bord+1:end-bord,bord+1:end-bord] + """ + stu = tu[:] + stv = tv[:] + su = u[:] + sv = v[:] + + idxUnknow = (abs(stu) > UNKNOWN_FLOW_THRESH) | (abs(stv) > UNKNOWN_FLOW_THRESH) + stu[idxUnknow] = 0 + stv[idxUnknow] = 0 + su[idxUnknow] = 0 + sv[idxUnknow] = 0 + + ind2 = [(np.absolute(stu) > smallflow) | (np.absolute(stv) > smallflow)] + index_su = su[ind2] + index_sv = sv[ind2] + an = 1.0 / np.sqrt(index_su ** 2 + index_sv ** 2 + 1) + un = index_su * an + vn = index_sv * an + + index_stu = stu[ind2] + index_stv = stv[ind2] + tn = 1.0 / np.sqrt(index_stu ** 2 + index_stv ** 2 + 1) + tun = index_stu * tn + tvn = index_stv * tn + + """ + angle = un * tun + vn * tvn + (an * tn) + index = [angle == 1.0] + angle[index] = 0.999 + ang = np.arccos(angle) + mang = np.mean(ang) + mang = mang * 180 / np.pi + """ + + epe = np.sqrt((stu - su) ** 2 + (stv - sv) ** 2) + epe = epe[ind2] + return np.mean(epe) + + +def flow_to_image(flow, display=False): + """ + Convert flow into middlebury color code image + :param flow: optical flow map + :return: optical flow image in middlebury color + """ + u = flow[:, :, 0] + v = flow[:, :, 1] + + maxu = -999.0 + maxv = -999.0 + minu = 999.0 + minv = 999.0 + + idxUnknow = (abs(u) > UNKNOWN_FLOW_THRESH) | (abs(v) > UNKNOWN_FLOW_THRESH) + u[idxUnknow] = 0 + v[idxUnknow] = 0 + + maxu = max(maxu, np.max(u)) + minu = min(minu, np.min(u)) + + maxv = max(maxv, np.max(v)) + minv = min(minv, np.min(v)) + + rad = np.sqrt(u ** 2 + v ** 2) + maxrad = max(-1, np.max(rad)) + + if display: + print("max flow: %.4f\nflow range:\nu = %.3f .. %.3f\nv = %.3f .. %.3f" % (maxrad, minu, maxu, minv, maxv)) + + u = u / (maxrad + np.finfo(float).eps) + v = v / (maxrad + np.finfo(float).eps) + + img = compute_color(u, v) + + idx = np.repeat(idxUnknow[:, :, np.newaxis], 3, axis=2) + img[idx] = 0 + + return np.uint8(img) + + +def evaluate_flow_file(gt, pred): + """ + evaluate the estimated optical flow end point error according to ground truth provided + :param gt: ground truth file path + :param pred: estimated optical flow file path + :return: end point error, float32 + """ + # Read flow files and calculate the errors + gt_flow = read_flow(gt) # ground truth flow + eva_flow = read_flow(pred) # predicted flow + # Calculate errors + average_pe = flow_error(gt_flow[:, :, 0], gt_flow[:, :, 1], eva_flow[:, :, 0], eva_flow[:, :, 1]) + return average_pe + + +def evaluate_flow(gt_flow, pred_flow): + """ + gt: ground-truth flow + pred: estimated flow + """ + return flow_error( + gt_flow[:, :, 0], + gt_flow[:, :, 1], + pred_flow[:, :, 0], + pred_flow[:, :, 1], + ) + + +""" +============== +Disparity Section +============== +""" + + +def read_disp_png(file_name): + """ + Read optical flow from KITTI .png file + :param file_name: name of the flow file + :return: optical flow data in matrix + """ + image_object = png.Reader(filename=file_name) + image_direct = image_object.asDirect() + image_data = list(image_direct[2]) + (w, h) = image_direct[3]["size"] + channel = len(image_data[0]) / w + flow = np.zeros((h, w, channel), dtype=np.uint16) + for i in range(len(image_data)): + for j in range(channel): + flow[i, :, j] = image_data[i][j::channel] + return flow[:, :, 0] / 256 + + +def disp_to_flowfile(disp, filename): + """ + Read KITTI disparity file in png format + :param disp: disparity matrix + :param filename: the flow file name to save + :return: None + """ + f = open(filename, "wb") + magic = np.array([202021.25], dtype=np.float32) + (height, width) = disp.shape[0:2] + w = np.array([width], dtype=np.int32) + h = np.array([height], dtype=np.int32) + empty_map = np.zeros((height, width), dtype=np.float32) + data = np.dstack((disp, empty_map)) + magic.tofile(f) + w.tofile(f) + h.tofile(f) + data.tofile(f) + f.close() + + +""" +============== +Image Section +============== +""" + + +def read_image(filename): + """ + Read normal image of any format + :param filename: name of the image file + :return: image data in matrix uint8 type + """ + img = Image.open(filename) + return np.array(img) + + +def warp_image(im, flow): + """ + Use optical flow to warp image to the next + :param im: image to warp + :param flow: optical flow + :return: warped image + """ + from scipy import interpolate + + image_height = im.shape[0] + image_width = im.shape[1] + flow_height = flow.shape[0] + flow_width = flow.shape[1] + n = image_height * image_width + (iy, ix) = np.float32(np.mgrid[0:image_height, 0:image_width]) + (fy, fx) = np.float32(np.mgrid[0:flow_height, 0:flow_width]) + fx += flow[:, :, 0].astype(np.float) + fy += flow[:, :, 1] + mask = np.logical_or(fx < 0, fx > flow_width) + mask = np.logical_or(mask, fy < 0) + mask = np.logical_or(mask, fy > flow_height) + fx = np.minimum(np.maximum(fx, 0), flow_width) + fy = np.minimum(np.maximum(fy, 0), flow_height) + points = np.concatenate((ix.reshape(n, 1), iy.reshape(n, 1)), axis=1) + xi = np.concatenate((fx.reshape(n, 1), fy.reshape(n, 1)), axis=1) + warp = np.zeros((image_height, image_width, im.shape[2])) + for i in range(im.shape[2]): + channel = im[:, :, i] + # plt.imshow(channel, cmap='gray') + values = channel.reshape(n, 1) + new_channel = interpolate.griddata(points, values, xi, method="cubic") + new_channel = np.reshape(new_channel, [flow_height, flow_width]) + new_channel[mask] = 1 + warp[:, :, i] = new_channel.astype(np.uint8) + + return warp.astype(np.uint8) + + +""" +============== +Others +============== +""" + + +def scale_image(image, new_range): + """ + Linearly scale the image into desired range + :param image: input image + :param new_range: the new range to be aligned + :return: image normalized in new range + """ + min_val = np.min(image).astype(np.float32) + max_val = np.max(image).astype(np.float32) + min_val_new = np.array(min(new_range), dtype=np.float32) + max_val_new = np.array(max(new_range), dtype=np.float32) + scaled_image = (image - min_val) / (max_val - min_val) * (max_val_new - min_val_new) + min_val_new + return scaled_image.astype(np.uint8) + + +def compute_color(u, v): + """ + compute optical flow color map + :param u: optical flow horizontal map + :param v: optical flow vertical map + :return: optical flow in color code + """ + [h, w] = u.shape + img = np.zeros([h, w, 3]) + nanIdx = np.isnan(u) | np.isnan(v) + u[nanIdx] = 0 + v[nanIdx] = 0 + + colorwheel = make_color_wheel() + ncols = np.size(colorwheel, 0) + + rad = np.sqrt(u ** 2 + v ** 2) + + a = np.arctan2(-v, -u) / np.pi + + fk = (a + 1) / 2 * (ncols - 1) + 1 + + k0 = np.floor(fk).astype(int) + + k1 = k0 + 1 + k1[k1 == ncols + 1] = 1 + f = fk - k0 + + for i in range(np.size(colorwheel, 1)): + tmp = colorwheel[:, i] + col0 = tmp[k0 - 1] / 255 + col1 = tmp[k1 - 1] / 255 + col = (1 - f) * col0 + f * col1 + + idx = rad <= 1 + col[idx] = 1 - rad[idx] * (1 - col[idx]) + notidx = np.logical_not(idx) + + col[notidx] *= 0.75 + img[:, :, i] = np.uint8(np.floor(255 * col * (1 - nanIdx))) + + return img + + +def make_color_wheel(): + """ + Generate color wheel according Middlebury color code + :return: Color wheel + """ + RY = 15 + YG = 6 + GC = 4 + CB = 11 + BM = 13 + MR = 6 + + ncols = RY + YG + GC + CB + BM + MR + + colorwheel = np.zeros([ncols, 3]) + + col = 0 + + # RY + colorwheel[0:RY, 0] = 255 + colorwheel[0:RY, 1] = np.transpose(np.floor(255 * np.arange(0, RY) / RY)) + col += RY + + # YG + colorwheel[col : col + YG, 0] = 255 - np.transpose(np.floor(255 * np.arange(0, YG) / YG)) + colorwheel[col : col + YG, 1] = 255 + col += YG + + # GC + colorwheel[col : col + GC, 1] = 255 + colorwheel[col : col + GC, 2] = np.transpose(np.floor(255 * np.arange(0, GC) / GC)) + col += GC + + # CB + colorwheel[col : col + CB, 1] = 255 - np.transpose(np.floor(255 * np.arange(0, CB) / CB)) + colorwheel[col : col + CB, 2] = 255 + col += CB + + # BM + colorwheel[col : col + BM, 2] = 255 + colorwheel[col : col + BM, 0] = np.transpose(np.floor(255 * np.arange(0, BM) / BM)) + col += +BM + + # MR + colorwheel[col : col + MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR)) + colorwheel[col : col + MR, 0] = 255 + + return colorwheel diff --git a/utils/tb_image_recorder.py b/utils/tb_image_recorder.py new file mode 100644 index 0000000..39ee0e4 --- /dev/null +++ b/utils/tb_image_recorder.py @@ -0,0 +1,28 @@ +import threading + +import numpy as np + + +class TBImageRecorder(threading.Thread): + """ + TBImageRecorder + """ + + def __init__(self, tb_writer, func, queue): + super(TBImageRecorder, self).__init__() + self._tb_writer = tb_writer + self._func = func + self._queue = queue + + def run(self): + while True: + msgs, iter_index = self._queue.get() + if msgs: + img_info = self._func(*msgs) + + print("logging the images") + for tag, images in img_info.items(): + if images is not None: + self._tb_writer.add_image(tag, np.clip(images, 0, 255).astype(np.uint8), iter_index) + else: + break diff --git a/utils/util_tensorboard.py b/utils/util_tensorboard.py new file mode 100644 index 0000000..a0e68cc --- /dev/null +++ b/utils/util_tensorboard.py @@ -0,0 +1,43 @@ +import threading + +import numpy as np + +from utils.util import to_np + + +def histogram_logger(tb_writer, iter_index, net=None): + if net is not None: + for tag, value in net.named_parameters(): + tag = tag.replace('.', '/') + tb_writer.add_histogram(tag, to_np(value), iter_index) + tb_writer.add_histogram(tag + '/grad', to_np(value.grad), iter_index) + + +def value_logger(tb_writer, iter_index, loss_info): + for tag, value in loss_info.items(): + tb_writer.add_scalar(tag, value, iter_index) + + +class TBImageRecorder(threading.Thread): + """ + TBImageRecorder + """ + + def __init__(self, tb_writer, func, queue): + super(TBImageRecorder, self).__init__() + self._tb_writer = tb_writer + self._func = func + self._queue = queue + + def run(self): + while True: + msgs, iter_index = self._queue.get() + if msgs: + img_info = self._func(*msgs) + + print("logging the images") + for tag, images in img_info.items(): + if images is not None: + self._tb_writer.add_image(tag, np.clip(images, 0, 255).astype(np.uint8), iter_index) + else: + break diff --git a/utils/warping.py b/utils/warping.py new file mode 100644 index 0000000..7df783d --- /dev/null +++ b/utils/warping.py @@ -0,0 +1,30 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def get_grid(x): + torchHorizontal = ( + torch.linspace(-1.0, 1.0, x.size(3)).view(1, 1, 1, x.size(3)).expand(x.size(0), 1, x.size(2), x.size(3)) + ) + torchVertical = ( + torch.linspace(-1.0, 1.0, x.size(2)).view(1, 1, x.size(2), 1).expand(x.size(0), 1, x.size(2), x.size(3)) + ) + return torch.cat([torchHorizontal, torchVertical], 1).cuda() + + +class WarpingLayer(nn.Module): + def __init__(self, device): + super(WarpingLayer, self).__init__() + self.device = device + + def forward(self, x, flow): + # WarpingLayer uses F.grid_sample, which expects normalized grid + # we still output unnormalized flow for the convenience of comparing EPEs with FlowNet2 and original code + # so here we need to denormalize the flow + flow_for_grip = torch.zeros_like(flow).cuda() + flow_for_grip[:, 0, :, :] = flow[:, 0, :, :] / ((flow.size(3) - 1.0) / 2.0) + flow_for_grip[:, 1, :, :] = flow[:, 1, :, :] / ((flow.size(2) - 1.0) / 2.0) + + grid = (get_grid(x) + flow_for_grip).permute(0, 2, 3, 1) + return F.grid_sample(x, grid, align_corners=True) diff --git a/warping.py b/warping.py new file mode 100644 index 0000000..8cdda79 --- /dev/null +++ b/warping.py @@ -0,0 +1,30 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def get_grid(x): + torchHorizontal = ( + torch.linspace(-1.0, 1.0, x.size(3)).view(1, 1, 1, x.size(3)).expand(x.size(0), 1, x.size(2), x.size(3)) + ) + torchVertical = ( + torch.linspace(-1.0, 1.0, x.size(2)).view(1, 1, x.size(2), 1).expand(x.size(0), 1, x.size(2), x.size(3)) + ) + return torch.cat([torchHorizontal, torchVertical], 1).cuda() + + +class WarpingLayer(nn.Module): + def __init__(self, device): + super(WarpingLayer, self).__init__() + self.device = device + + def forward(self, x, flow): + # WarpingLayer uses F.VGG19_pytorch, which expects normalized grid + # we still output unnormalized flow for the convenience of comparing EPEs with FlowNet2 and original code + # so here we need to denormalize the flow + flow_for_grip = torch.zeros_like(flow).cuda() + flow_for_grip[:, 0, :, :] = flow[:, 0, :, :] / ((flow.size(3) - 1.0) / 2.0) + flow_for_grip[:, 1, :, :] = flow[:, 1, :, :] / ((flow.size(2) - 1.0) / 2.0) + + grid = (get_grid(x) + flow_for_grip).permute(0, 2, 3, 1) + return F.VGG19_pytorch(x, grid, align_corners=True)