File size: 4,531 Bytes
f12ab4c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
import numpy as np
import torch
#----------------------------------------------------------------------------
# Projection and transformation matrix helpers.
#----------------------------------------------------------------------------
def projection(x=0.1, n=1.0, f=50.0):
return np.array([[n/x, 0, 0, 0],
[ 0, n/-x, 0, 0],
[ 0, 0, -(f+n)/(f-n), -(2*f*n)/(f-n)],
[ 0, 0, -1, 0]]).astype(np.float32)
def translate(x, y, z):
return np.array([[1, 0, 0, x],
[0, 1, 0, y],
[0, 0, 1, z],
[0, 0, 0, 1]]).astype(np.float32)
def rotate_x(a):
s, c = np.sin(a), np.cos(a)
return np.array([[1, 0, 0, 0],
[0, c, s, 0],
[0, -s, c, 0],
[0, 0, 0, 1]]).astype(np.float32)
def rotate_y(a):
s, c = np.sin(a), np.cos(a)
return np.array([[ c, 0, s, 0],
[ 0, 1, 0, 0],
[-s, 0, c, 0],
[ 0, 0, 0, 1]]).astype(np.float32)
def random_rotation_translation(t):
m = np.random.normal(size=[3, 3])
m[1] = np.cross(m[0], m[2])
m[2] = np.cross(m[0], m[1])
m = m / np.linalg.norm(m, axis=1, keepdims=True)
m = np.pad(m, [[0, 1], [0, 1]], mode='constant')
m[3, 3] = 1.0
m[:3, 3] = np.random.uniform(-t, t, size=[3])
return m
#----------------------------------------------------------------------------
# Bilinear downsample by 2x.
#----------------------------------------------------------------------------
def bilinear_downsample(x):
w = torch.tensor([[1, 3, 3, 1], [3, 9, 9, 3], [3, 9, 9, 3], [1, 3, 3, 1]], dtype=torch.float32, device=x.device) / 64.0
w = w.expand(x.shape[-1], 1, 4, 4)
x = torch.nn.functional.conv2d(x.permute(0, 3, 1, 2), w, padding=1, stride=2, groups=x.shape[-1])
return x.permute(0, 2, 3, 1)
#----------------------------------------------------------------------------
# Image display function using OpenGL.
#----------------------------------------------------------------------------
_glfw_window = None
def display_image(image, zoom=None, size=None, title=None): # HWC
# Import OpenGL and glfw.
import OpenGL.GL as gl
import glfw
# Zoom image if requested.
image = np.asarray(image)
if size is not None:
assert zoom is None
zoom = max(1, size // image.shape[0])
if zoom is not None:
image = image.repeat(zoom, axis=0).repeat(zoom, axis=1)
height, width, channels = image.shape
# Initialize window.
if title is None:
title = 'Debug window'
global _glfw_window
if _glfw_window is None:
glfw.init()
_glfw_window = glfw.create_window(width, height, title, None, None)
glfw.make_context_current(_glfw_window)
glfw.show_window(_glfw_window)
glfw.swap_interval(0)
else:
glfw.make_context_current(_glfw_window)
glfw.set_window_title(_glfw_window, title)
glfw.set_window_size(_glfw_window, width, height)
# Update window.
glfw.poll_events()
gl.glClearColor(0, 0, 0, 1)
gl.glClear(gl.GL_COLOR_BUFFER_BIT)
gl.glWindowPos2f(0, 0)
gl.glPixelStorei(gl.GL_UNPACK_ALIGNMENT, 1)
gl_format = {3: gl.GL_RGB, 2: gl.GL_RG, 1: gl.GL_LUMINANCE}[channels]
gl_dtype = {'uint8': gl.GL_UNSIGNED_BYTE, 'float32': gl.GL_FLOAT}[image.dtype.name]
gl.glDrawPixels(width, height, gl_format, gl_dtype, image[::-1])
glfw.swap_buffers(_glfw_window)
if glfw.window_should_close(_glfw_window):
return False
return True
#----------------------------------------------------------------------------
# Image save helper.
#----------------------------------------------------------------------------
def save_image(fn, x):
import imageio
x = np.rint(x * 255.0)
x = np.clip(x, 0, 255).astype(np.uint8)
imageio.imsave(fn, x)
#----------------------------------------------------------------------------
|