class ImageSubscriber():
def __init__(self):
self.cvbridge = CvBridge()
self.img_buf = None
self.image_subs = rospy.Subscriber("image", Image, callback=self._callback, queue_size=1)
self.image_viewer = ImageViewer()
self.test = rospy.get_param("/stop_param")
def wait_for_image(self):
if self.test is True:
self.img_buf = self.image_viewer.get_test_camera()
if self.img_buf is None:
return None, None
key = self.image_viewer.show(self.img_buf)
img = self.img_buf
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, dsize=(HEIGHT, WIDTH))
img = (img.astype(np.float32) - 128) / 256
self.img_buf = None
return img, key
def _callback(self, msg):
cv_img = self.cvbridge.imgmsg_to_cv2(msg)
self.img_buf = cv_img
class GameState:
def __init__(self, screen_type="console"):
self.game = core.Game()
self.state = core.state
self.screen_type = screen_type
self.console = curses.initscr()
self.score = 0
self.viewer = None
self.screen = None
pygame.init()
def initial_state(self):
return (
helpers.translate_state(self.state, self.game.table),
0.5,
False,
)
def frame_step(self, action):
return self.do_action(action)
def print_console(self, board):
self.console.addstr(0, 0, board)
self.console.refresh()
def render(self, mode="human"):
if self.viewer is None:
self.viewer = ImageViewer(caption="SnakeEnv", height=500, width=500,)
# show the screen on the image viewer
self.viewer.show(self.screen)
def find_loc(self, row_ix, col_ix):
loc = [
(col_ix * 4.2),
(row_ix * 4.2),
(col_ix * 4.2 + 4.2),
(row_ix * 4.2 + 4.2),
]
return loc
def draw(self, data):
img = Image.new("RGB", (128, 128))
d = ImageDraw.Draw(img)
for row_ix, row in enumerate(data):
for col_ix, col in enumerate(row):
if col == "#":
d.rectangle(
self.find_loc(row_ix, col_ix), fill=(255, 0, 0),
)
elif col == "X":
d.rectangle(
self.find_loc(row_ix, col_ix), fill=(0, 255, 0),
)
elif col == "O":
d.ellipse(
self.find_loc(row_ix, col_ix), fill=(0, 0, 255),
)
elif col == "H":
d.rectangle(
self.find_loc(row_ix, col_ix), fill=(128, 0, 128),
)
return np.array(img)
def print_console_data(self, data):
self.console.addstr(30, 0, data)
self.console.refresh()
def do_action(self, action):
"""
action: In range [0, 1, 2]
0 do nothing
1 left
2 right
1, 2 will be translated to respective value according to the snake traveling path.
"""
self.screen = self.draw(self.game.table)
new_state, eaten, is_state_terminal = self.driveSnake(action, self.state)
self.state = new_state
if eaten:
reward = 1
self.score += 1
elif is_state_terminal and not eaten:
reward = -1
self.score = 0
else:
reward = -0.05
return (
helpers.translate_state(new_state, self.game.table),
reward,
is_state_terminal,
self.score,
)
def driveSnake(self, direction, state):
traveling_path = state["moves"][0]
if traveling_path == core.NORTH:
move = core.WEST if direction == 1 else core.EAST
elif traveling_path == core.SOUTH:
move = core.EAST if direction == 1 else core.WEST
elif traveling_path == core.WEST:
move = core.SOUTH if direction == 1 else core.NORTH
elif traveling_path == core.EAST:
move = core.NORTH if direction == 1 else core.SOUTH
else:
raise ValueError("Invalid travelling path.")
if direction == 0:
move = self.state["moves"][0]
new_state = core.enqeue_move(state, move)
new_state, eaten, is_state_terminal = core.next_state(new_state)
self.game.make(new_state, new_state["apple"])
if self.screen_type == "console":
board = str(self.game)
self.print_console(board)
else:
self.render() # use a gui
return new_state, eaten, is_state_terminal
kwargs = evaluation['kwargs']
test_perc = evaluation['test_perc']
all_dices.append(dice_vals)
all_filenames = filenames if all_filenames is None else all_filenames
pairs = list(zip(dice_vals, filenames))
pairs.sort()
dice_vals = [x for x, _ in pairs]
filenames = [x for _, x in pairs]
plt.plot(dice_vals, label=f'dice {i}')
plt.axhline(y=mean, color='rgb'[i % 3], ls='--', label=f'mean {i}')
print(
f'{FILENAME:<32} mean={mean:.04f} time={time:.02f}s fps={len(filenames)/time:.02f} test_perc={test_perc} kwargs={kwargs}'
)
plt.title('\n'.join(show_filenames))
plt.legend()
plt.show()
if len(all_dices) == 2:
worst = [(max(a - b, b - a), f)
for a, b, f in zip(*all_dices, filenames)]
worst.sort()
iv = ImageViewer()
for s, f in worst[-12:]:
iv.add(cv2.imread(f'data_test/paintings_gt/imgs/{f}'),
cmap='bgr',
title=f'{s}')
iv.show()
