diff options
Diffstat (limited to 'util')
| -rw-r--r-- | util/png.py | 45 | ||||
| -rw-r--r-- | util/util.py | 19 | ||||
| -rw-r--r-- | util/visualizer.py | 12 |
3 files changed, 35 insertions, 41 deletions
diff --git a/util/png.py b/util/png.py index 0936cf0..3a750c0 100644 --- a/util/png.py +++ b/util/png.py @@ -1,32 +1,33 @@ import struct import zlib + def encode(buf, width, height): - """ buf: must be bytes or a bytearray in py3, a regular string in py2. formatted RGBRGB... """ - assert (width * height * 3 == len(buf)) - bpp = 3 + """ buf: must be bytes or a bytearray in py3, a regular string in py2. formatted RGBRGB... """ + assert (width * height * 3 == len(buf)) + bpp = 3 - def raw_data(): - # reverse the vertical line order and add null bytes at the start - row_bytes = width * bpp - for row_start in range((height - 1) * width * bpp, -1, -row_bytes): - yield b'\x00' - yield buf[row_start:row_start + row_bytes] + def raw_data(): + # reverse the vertical line order and add null bytes at the start + row_bytes = width * bpp + for row_start in range((height - 1) * width * bpp, -1, -row_bytes): + yield b'\x00' + yield buf[row_start:row_start + row_bytes] - def chunk(tag, data): - return [ - struct.pack("!I", len(data)), - tag, - data, - struct.pack("!I", 0xFFFFFFFF & zlib.crc32(data, zlib.crc32(tag))) - ] + def chunk(tag, data): + return [ + struct.pack("!I", len(data)), + tag, + data, + struct.pack("!I", 0xFFFFFFFF & zlib.crc32(data, zlib.crc32(tag))) + ] - SIGNATURE = b'\x89PNG\r\n\x1a\n' - COLOR_TYPE_RGB = 2 - COLOR_TYPE_RGBA = 6 - bit_depth = 8 - return b''.join( - [ SIGNATURE ] + + SIGNATURE = b'\x89PNG\r\n\x1a\n' + COLOR_TYPE_RGB = 2 + COLOR_TYPE_RGBA = 6 + bit_depth = 8 + return b''.join( + [SIGNATURE] + chunk(b'IHDR', struct.pack("!2I5B", width, height, bit_depth, COLOR_TYPE_RGB, 0, 0, 0)) + chunk(b'IDAT', zlib.compress(b''.join(raw_data()), 9)) + chunk(b'IEND', b'') diff --git a/util/util.py b/util/util.py index 4de0a74..26b259a 100644 --- a/util/util.py +++ b/util/util.py @@ -2,11 +2,13 @@ from __future__ import print_function import torch import numpy as np from PIL import Image -import inspect, re +import inspect +import re import numpy as np import os import collections + # Converts a Tensor into a Numpy array # |imtype|: the desired type of the converted numpy array def tensor2im(image_tensor, imtype=np.uint8): @@ -34,21 +36,6 @@ def save_image(image_numpy, image_path): image_pil = Image.fromarray(image_numpy) image_pil.save(image_path) -def info(object, spacing=10, collapse=1): - """Print methods and doc strings. - Takes module, class, list, dictionary, or string.""" - methodList = [e for e in dir(object) if isinstance(getattr(object, e), collections.Callable)] - processFunc = collapse and (lambda s: " ".join(s.split())) or (lambda s: s) - print( "\n".join(["%s %s" % - (method.ljust(spacing), - processFunc(str(getattr(object, method).__doc__))) - for method in methodList]) ) - -def varname(p): - for line in inspect.getframeinfo(inspect.currentframe().f_back)[3]: - m = re.search(r'\bvarname\s*\(\s*([A-Za-z_][A-Za-z0-9_]*)\s*\)', line) - if m: - return m.group(1) def print_numpy(x, val=True, shp=False): x = x.astype(np.float64) diff --git a/util/visualizer.py b/util/visualizer.py index e6e7cba..8bec8df 100644 --- a/util/visualizer.py +++ b/util/visualizer.py @@ -4,6 +4,7 @@ import ntpath import time from . import util from . import html +from scipy.misc import imresize class Visualizer(): @@ -123,7 +124,7 @@ class Visualizer(): log_file.write('%s\n' % message) # save image to the disk - def save_images(self, webpage, visuals, image_path): + def save_images(self, webpage, visuals, image_path, aspect_ratio=1.0): image_dir = webpage.get_image_dir() short_path = ntpath.basename(image_path[0]) name = os.path.splitext(short_path)[0] @@ -133,10 +134,15 @@ class Visualizer(): txts = [] links = [] - for label, image_numpy in visuals.items(): + for label, im in visuals.items(): image_name = '%s_%s.png' % (name, label) save_path = os.path.join(image_dir, image_name) - util.save_image(image_numpy, save_path) + h, w, _ = im.shape + if aspect_ratio > 1.0: + im = imresize(im, (h, int(w * aspect_ratio)), interp='bicubic') + if aspect_ratio < 1.0: + im = imresize(im, (int(h / aspect_ratio), w), interp='bicubic') + util.save_image(im, save_path) ims.append(image_name) txts.append(label) |
