def show_image(im, cmap=None):
fig = plt.figure(figsize=(3,3))
plt.imshow(im, cmap=cmap)
plt.grid(False)
plt.show()
def gen_img(start, target, fill=1, img_size=10):
# Generates empty image
img = np.zeros((img_size, img_size), dtype=float)
start_row, start_col = None, None
if start > 0:
start_row = start
else:
start_col = np.abs(start)
if target == 0:
if start_row is None:
img[:, start_col] = fill
else:
img[start_row, :] = fill
else:
if start_col == 0:
start_col = 1
if target == 1:
if start_row is not None:
up = (range(start_row, -1, -1),
range(0, start_row + 1))
else:
up = (range(img_size - 1, start_col - 1, -1),
range(start_col, img_size))
img[up] = fill
else:
if start_row is not None:
down = (range(start_row, img_size, 1),
range(0, img_size - start_row))
else:
down = (range(0, img_size - 1 - start_col + 1),
range(start_col, img_size))
img[down] = fill
return 255 * img.reshape(1, img_size, img_size)
def generate_dataset(img_size=10, n_images=100, binary=True, seed=17):
np.random.seed(seed)
starts = np.random.randint(-(img_size - 1), img_size, size=(n_images,))
targets = np.random.randint(0, 3, size=(n_images,))
images = np.array([gen_img(s, t, img_size=img_size)
for s, t in zip(starts, targets)], dtype=np.uint8)
if binary:
targets = (targets > 0).astype(int)
return images, targets
def plot_images(images, targets, n_plot=30):
n_rows = n_plot // 6 + ((n_plot % 6) > 0)
fig, axes = plt.subplots(n_rows, 6, figsize=(9, 1.5 * n_rows))
axes = np.atleast_2d(axes)
for i, (image, target) in enumerate(zip(images[:n_plot], targets[:n_plot])):
row, col = i // 6, i % 6
ax = axes[row, col]
ax.set_title('#{} - Label:{}'.format(i, target), {'size': 12})
# plot filter channel in grayscale
ax.imshow(image.squeeze(), cmap='gray', vmin=0, vmax=1)
for ax in axes.flat:
ax.set_xticks([])
ax.set_yticks([])
ax.label_outer()
plt.tight_layout()
return fig
class TransformedTensorDataset(Dataset):
def __init__(self, x, y, transform=None):
self.x = x
self.y = y
self.transform = transform
def __getitem__(self, index):
x = self.x[index]
if self.transform:
x = self.transform(x)
return x, self.y[index]
def __len__(self):
return len(self.x)