mirror of
https://github.com/wassname/CoordConv.git
synced 2026-08-20 12:00:37 +08:00
91 lines
2.6 KiB
Python
91 lines
2.6 KiB
Python
import os
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
from sklearn.model_selection import train_test_split
|
|
|
|
datatype = 'uniform'
|
|
assert datatype in ['uniform', 'quadrant']
|
|
|
|
if not os.path.exists('data-uniform/'):
|
|
os.makedirs('data-uniform/')
|
|
|
|
if not os.path.exists('data-quadrant/'):
|
|
os.makedirs('data-quadrant/')
|
|
|
|
np.random.seed(0)
|
|
torch.manual_seed(0)
|
|
|
|
onehots = np.pad(np.eye(3136, dtype='float32').reshape((3136, 56, 56, 1)),
|
|
((0, 0), (4, 4), (4, 4), (0, 0)), mode="constant")
|
|
onehots = onehots.transpose(0, 3, 1, 2)
|
|
|
|
onehots_tensor = torch.from_numpy(onehots)
|
|
|
|
conv_layer = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=(9, 9), padding=4, stride=1)
|
|
w = torch.ones(1, 1, 9, 9)
|
|
conv_layer.weight.data = w
|
|
|
|
images_tensor = conv_layer(onehots_tensor)
|
|
images = images_tensor.detach().numpy()
|
|
|
|
if datatype == 'uniform':
|
|
# Create the uniform datasets
|
|
indices = np.arange(0, len(onehots), dtype='int32')
|
|
train, test = train_test_split(indices, test_size=0.2, random_state=0)
|
|
|
|
train_onehot = onehots[train]
|
|
train_images = images[train]
|
|
|
|
test_onehot = onehots[test]
|
|
test_images = images[test]
|
|
|
|
np.save('data-uniform/train_onehot.npy', train_onehot)
|
|
np.save('data-uniform/train_images.npy', train_images)
|
|
np.save('data-uniform/test_onehot.npy', test_onehot)
|
|
np.save('data-uniform/test_images.npy', test_images)
|
|
else:
|
|
pos_quadrant = np.where(onehots == 1.0)
|
|
# print(onehots.shape)
|
|
X = pos_quadrant[2]
|
|
Y = pos_quadrant[3]
|
|
|
|
train_set = []
|
|
test_set = []
|
|
|
|
train_ids = []
|
|
test_ids = []
|
|
|
|
for i, (x, y) in enumerate(zip(X, Y)):
|
|
if x > 32 and y > 32: # 4th quadrant
|
|
test_ids.append(i)
|
|
test_set.append([x, y])
|
|
else:
|
|
train_ids.append(i)
|
|
train_set.append([x, y])
|
|
|
|
train_set = np.array(train_set)
|
|
test_set = np.array(test_set)
|
|
|
|
train_set = train_set[:, :, None, None]
|
|
test_set = test_set[:, :, None, None]
|
|
|
|
print(train_set.shape)
|
|
print(test_set.shape)
|
|
|
|
train_onehot = onehots[train_ids]
|
|
test_onehot = onehots[test_ids]
|
|
|
|
train_images = images[train_ids]
|
|
test_images = images[test_ids]
|
|
|
|
print(train_onehot.shape, test_onehot.shape)
|
|
print(train_images.shape, test_images.shape)
|
|
|
|
np.save('data-quadrant/train_set.npy', train_set)
|
|
np.save('data-quadrant/test_set.npy', test_set)
|
|
np.save('data-quadrant/train_onehot.npy', train_onehot)
|
|
np.save('data-quadrant/train_images.npy', train_images)
|
|
np.save('data-quadrant/test_onehot.npy', test_onehot)
|
|
np.save('data-quadrant/test_images.npy', test_images)
|