-
Notifications
You must be signed in to change notification settings - Fork 71
/
render.py
68 lines (53 loc) · 1.89 KB
/
render.py
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
import math
import pygame
i = 0
def multi_view_render(images, unit_dimension, actor_configs):
"""Render images based on pygame > 1.9.4
Args:
images (dict):
unit_dimension (list): window size, e.g., [84, 84]
actor_configs (dict): configs of actors
Returns:
N/A.
"""
global i
pygame.init()
surface_seq = ()
poses, window_dim = get_surface_poses(
len(images), unit_dimension, images.keys())
# Get all surfaces.
for actor_id, im in images.items():
if not actor_configs[actor_id]["render"]:
continue
surface = pygame.surfarray.make_surface(im.swapaxes(0, 1) * 128 + 128)
surface_seq += ((surface, (poses[actor_id][1], poses[actor_id][0])), )
display = pygame.display.set_mode((window_dim[0], window_dim[1]),
pygame.HWSURFACE | pygame.DOUBLEBUF)
display.blits(blit_sequence=surface_seq, doreturn=1)
pygame.display.flip()
# save to disk
# pygame.image.save(display,
# "/mnt/DATADRIVE1/pygame_surfs/" + str(i) + ".jpeg")
i += 1
def get_surface_poses(subwindow_num, unit_dimension, ids):
"""Calculate the poses of sub-windows of actors
Args:
subwindow_num (int): number of sub-windows(actors)
window_dimension (list): E.g., [800, 600]
ids (list): list of actor ids.
Returns:
dict: return position dicts in pygame window (start from left corner).
E.g., {"vehiclie":[0,0]}
"""
num = subwindow_num
unit_x = unit_dimension[0]
unit_y = unit_dimension[1]
row_num = math.ceil(math.sqrt(num))
max_x = row_num * unit_x
max_y = row_num * unit_y
poses = {}
for i, id in enumerate(ids):
x_pos = math.floor(i / row_num) * unit_x
y_pos = math.floor(i % row_num) * unit_y
poses[id] = [x_pos, y_pos]
return poses, [max_x, max_y]