File size: 8,628 Bytes
954caab |
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 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 |
import numpy as np
from PIL import Image
import torch
from einops import einsum, rearrange
from .permutations import make_jigsaw_perm, get_inv_perm
from .view_permute import PermuteView
from .jigsaw_helpers import get_jigsaw_pieces
class JigsawView(PermuteView):
'''
Implements a 4x4 jigsaw puzzle view...
'''
def __init__(self, seed=11):
'''
'''
# Get pixel permutations, corresponding to jigsaw permutations
self.perm_64, _ = make_jigsaw_perm(64, seed=seed)
self.perm_256, (jigsaw_perm) = make_jigsaw_perm(256, seed=seed)
# keep track of jigsaw permutation as well
self.piece_perms, self.edge_swaps = jigsaw_perm
# Init parent PermuteView, with above pixel perms
super().__init__(self.perm_64, self.perm_256)
def extract_pieces(self, im):
'''
Given an image, extract jigsaw puzzle pieces from it
im (PIL.Image) :
PIL Image of the jigsaw illusion
'''
im = np.array(im)
size = im.shape[0]
pieces = []
# Get jigsaw pieces
piece_masks = get_jigsaw_pieces(size)
# Save pieces
for piece_mask in piece_masks:
# Add mask as alpha mask to image
im_piece = np.concatenate([im, piece_mask[:,:,None] * 255], axis=2)
# Get extents of piece, and crop
x_min = np.nonzero(im_piece[:,:,-1].sum(0))[0].min()
x_max = np.nonzero(im_piece[:,:,-1].sum(0))[0].max()
y_min = np.nonzero(im_piece[:,:,-1].sum(1))[0].min()
y_max = np.nonzero(im_piece[:,:,-1].sum(1))[0].max()
im_piece = im_piece[y_min:y_max+1, x_min:x_max+1]
pieces.append(Image.fromarray(im_piece))
return pieces
def paste_piece(self, piece, x, y, theta, xc, yc, canvas_size=384):
'''
Given a PIL Image of a piece, place it so that it's center is at
(x,y) and it's rotate about that center at theta degrees
x (float) : x coordinate to place piece at
y (float) : y coordinate to place piece at
theta (float) : degrees to rotate piece about center
xc (float) : x coordinate of center of piece
yc (float) : y coordinate of center of piece
'''
# Make canvas
canvas = Image.new("RGBA",
(canvas_size, canvas_size),
(255, 255, 255, 0))
# Past piece so center is at (x, y)
canvas.paste(piece, (x-xc,y-yc), piece)
# Rotate about (x, y)
canvas = canvas.rotate(theta, resample=Image.BILINEAR, center=(x, y))
return canvas
def make_frame(self, im, t, canvas_size=384, knot_seed=0):
'''
This function returns a PIL image of a frame animating a jigsaw
permutation. Pieces move and rotate from the identity view
(t = 0) to the rearranged view (t = 1) along splines.
The approach is as follows:
1. Extract all 16 pieces
2. Figure out start locations for each of these pieces (t=0)
3. Figure out how these pieces permute
4. Using these permutations, figure out end locations (t=1)
5. Make knots for splines, randomly offset normally from the
midpoint of the start and end locations
6. Paste pieces into correct locations, determined by
spline interpolation
im (PIL.Image) :
PIL image representing the jigsaw illusion
t (float) :
Interpolation parameter in [0,1] indicating what frame of the
animation to generate
canvas_size (int) :
Side length of the frame
knot_seed (int) :
Seed for random offsets for the knots
'''
im_size = im.size[0]
# Extract 16 jigsaw pieces
pieces = self.extract_pieces(im)
# Rotate all pieces to "base" piece orientation
pieces = [p.rotate(90 * (i % 4),
resample=Image.BILINEAR,
expand=1) for i, p in enumerate(pieces)]
# Get (hardcoded) start locations for each base piece, on a
# 4x4 grid centered on the origin.
corner_start_loc = np.array([-1.5, -1.5])
inner_start_loc = np.array([-0.5, -0.5])
edge_e_start_loc = np.array([-1.5, -0.5])
edge_f_start_loc = np.array([-1.5, 0.5])
base_start_locs = np.stack([corner_start_loc,
inner_start_loc,
edge_e_start_loc,
edge_f_start_loc])
# Construct all start locations by rotating around (0,0)
# by 90 degrees, 4 times, and concatenating the results
rot_mats = []
for theta in -np.arange(4) * 90 / 180 * np.pi:
rot_mat = np.array([[np.cos(theta), -np.sin(theta)],
[np.sin(theta), np.cos(theta)]])
rot_mats.append(rot_mat)
rot_mats = np.stack(rot_mats)
start_locs = einsum(base_start_locs, rot_mats,
'start i, rot j i -> start rot j')
start_locs = rearrange(start_locs,
'start rot j -> (start rot) j')
# Add rotation information to start locations
thetas = np.tile(np.arange(4) * -90, 4)[:, None]
start_locs = np.concatenate([start_locs, thetas], axis=1)
# Get explicit permutation of pieces from permutation metadata
perm = self.piece_perms + np.repeat(np.arange(4), 4) * 4
for edge_idx, to_swap in enumerate(self.edge_swaps):
if to_swap:
# Make swap permutation array
swap_perm = np.arange(16)
swap_perm[8 + edge_idx], swap_perm[12 + edge_idx] = \
swap_perm[12 + edge_idx], swap_perm[8 + edge_idx]
# Apply swap permutation after perm
perm = np.array([swap_perm[perm[i]] for i in range(16)])
# Get inverse perm (the actual permutation needed)...
perm_inv = get_inv_perm(torch.tensor(perm))
# ...and use it to get the final locations of pieces
end_locs = start_locs[perm_inv]
# Convert start and end locations to pixel coordinate system
start_locs[:,:2] = (start_locs[:,:2] + 2) * 64
end_locs[:,:2] = (end_locs[:,:2] + 2) * 64
# Add offset so pieces are centered on canvas
start_locs[:,:2] = start_locs[:,:2] + (canvas_size - im_size) // 2
end_locs[:,:2] = end_locs[:,:2] + (canvas_size - im_size) // 2
# Get random offsets from middle for spline knot (so path is pretty)
# Wrapped in a set seed
original_state = np.random.get_state()
np.random.seed(knot_seed)
rand_offsets = np.random.rand(16, 1) * 2 - 1
rand_offsets = rand_offsets * 2
eps = np.random.randn(16, 2) # Add epsilon for divide by zero
np.random.set_state(original_state)
# Make spline knots by taking average of start and end,
# and offsetting by some amount normal from the line
avg_locs = (start_locs[:, :2] + end_locs[:, :2]) / 2.
norm = (end_locs[:, :2] - start_locs[:, :2])
norm = norm + eps
norm = norm / np.linalg.norm(norm, axis=1, keepdims=True)
rot_mat = np.array([[0,1], [-1,0]])
norm = norm @ rot_mat
rand_offsets = rand_offsets * (im_size / 4)
knot_locs = avg_locs + norm * rand_offsets
# Paste pieces on to a canvas
canvas = Image.new("RGBA", (canvas_size, canvas_size), (255,255,255,255))
for i in range(16):
# Get start and end coords
y_0, x_0, theta_0 = start_locs[i]
y_1, x_1, theta_1 = end_locs[i]
y_k, x_k = knot_locs[i]
# Take spline interpolation for x and y
x_int_0 = x_0 * (1-t) + x_k * t
y_int_0 = y_0 * (1-t) + y_k * t
x_int_1 = x_k * (1-t) + x_1 * t
y_int_1 = y_k * (1-t) + y_1 * t
x = int(np.round(x_int_0 * (1-t) + x_int_1 * t))
y = int(np.round(y_int_0 * (1-t) + y_int_1 * t))
# Just take normal interpolation for theta
theta = int(np.round(theta_0 * (1-t) + theta_1 * t))
# Get piece in location and rotation
xc = yc = im_size // 4 // 2
pasted_piece = self.paste_piece(pieces[i], x, y, theta, xc, yc)
canvas.paste(pasted_piece, (0,0), pasted_piece)
return canvas
|