def test_flowfield():
u = UniImageViewer()
x = bad_monkey()
theta = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]).expand(1, -1, -1)
grid = F.affine_grid(theta, x.shape)
out = F.grid_sample(x, grid)
u.render(out[0], block=True)
def test_pong():
v = UniImageViewer()
l = UniImageViewer(title='processed', screen_resolution=(32, 32))
env = gym.make('Pong-v0')
s = env.reset()
done = False
while not done:
s, r, done, info = env.step(cma_es.sample())
v.render(s)
s = d.pong_color_prepro(s)
#s = cv2.cvtColor(s, cv2.COLOR_RGB2GRAY)
#s = s[34:168, :]
#s = cv2.resize(s, dsize=(32, 32), interpolation=cv2.INTER_AREA)
l.render(s)
def evaluate(args, env, policy, render=False):
with torch.no_grad():
def get_action(s, prepro, transform, view, policy, action_map, device):
s = prepro(s)
s_t = transform(s).unsqueeze(0).type(policy.dtype).to(device)
kp = view(s_t)
p = softmax(keypoints.models.knn.flatten().matmul(policy), dim=0)
a = Categorical(p).sample()
a = action_map(a)
return a, kp
v = UniImageViewer()
datapack = ds.datasets[args.dataset]
if args.model_type != 'nop':
transporter_net = transporter.make(args, map_device='cpu')
view = Keypoints(transporter_net).to(args.device)
else:
view = nop
s = env.reset()
a, kp = get_action(s, datapack.prepro, datapack.transforms, view,
policy, datapack.action_map, args.device)
done = False
reward = 0.0
while not done:
s, r, done, i = env.step(a)
reward += r
a, kp = get_action(s, datapack.prepro, datapack.transforms, view,
policy, datapack.action_map, args.device)
if render:
if args.model_keypoints:
s = datapack.prepro(s)
s = TVF.to_tensor(s).unsqueeze(0)
s = plot_keypoints_on_image(kp[0], s[0])
v.render(s)
else:
env.render()
return reward
def test_tps_random():
images = []
u = UniImageViewer(screen_resolution=(2400, 1200))
x = bad_monkey()
for i in range(5, 10):
set = []
for _ in range(8):
pass
# set.append(tps_random(x, num_control_points=20, var=1 / i))
st = torch.cat(set, dim=2)
images.append(st)
img = torch.cat(images, dim=3)
u.render(img, block=True)
def test_tps():
u = UniImageViewer()
x = bad_monkey()
theta = torch.tensor([[[0.0, 0.0], [0., 0.], [0., 0.], [0., 0.], [0., 0.],
[0., 0.], [0.0, 0.0]]])
c = torch.tensor([
[0., 0],
[1., 0],
[1., 1],
[0, 1],
]).unsqueeze(0)
grid = tps_grid(theta, c, x.shape)
out = F.grid_sample(x, grid)
u.render(out[0], block=True)
def test_pong_fill():
l = 6000
display = True
ds = AtariDataset('Pong-v0', l, d.pong_prepro,
transforms = T.Compose([T.ToTensor(), T.ToPILImage()]),
end_trajectory=if_done_or_nonzero_reward)
assert len(ds) >= l
print(len(ds))
disp = UniImageViewer(screen_resolution=(512, 512))
for img1, img2 in ds:
#i = np.concatenate((img[0], img[1]), axis=1)
if display:
disp.render(img1)
time.sleep(0.03)
else:
plt.imshow(img1, cmap='gray', vmin=0, vmax=256.0)
plt.show()
from torchvision.transforms import functional as TVF
if __name__ == '__main__':
args = config.config()
with torch.no_grad():
v = UniImageViewer()
datapack = ds.datasets[args.dataset]
transporter_net = transporter.make(args).to(args.device)
if args.load is not None:
transporter_net.load(args.load)
env = gym.make(datapack.env)
while True:
s = env.reset()
done = False
while not done:
s, r, done, i = env.step(cma_es.sample())
s = datapack.prepro(s)
s_t = datapack.transforms(s).unsqueeze(0).to(args.device)
heatmap = transporter_net.keypoint(s_t)
kp = KF.spacial_logsoftmax(heatmap)
s = TVF.to_tensor(s).unsqueeze(0)
s = plot_keypoints_on_image(kp[0], s[0])
v.render(s)
time.sleep(0.04)
projection = pyrr.matrix44.create_orthogonal_projection_matrix(
*size.pyrr_orthogonal_proj)
glUniformMatrix4fv(proj_loc, 1, GL_FALSE, projection)
glUniformMatrix4fv(model_loc, 1, GL_FALSE, model)
glDrawElements(GL_TRIANGLES, offsets['screen'].len, GL_UNSIGNED_INT,
offsets['screen'].offset)
buffer = glReadPixels(anchr.x,
anchr.y,
size.width,
size.height,
format=GL_RGB,
type=GL_UNSIGNED_BYTE)
image = Image.frombuffer(mode='RGB',
size=(size.width, size.height),
data=buffer)
image.save(
f'/home/duane/PycharmProjects/keypoints/data/patches/pong/unclassified/pong_{imageid}.png'
)
pixel_array = np.frombuffer(buffer, dtype=np.uint8).reshape(
size.width, size.height, 3)
viewer.render(pixel_array)
imageid += 1
glfw.swap_buffers(window)
sleep(0.01)
# terminate glfw, free up allocated resources
glfw.terminate()