def save_image(self, log, step):
for i in range(self.env_batch):
if self.env.imgid[i] <= 10:
canvas = cv2.cvtColor(
(to_numpy(self.env.canvas[i].permute(1, 2, 0))),
cv2.COLOR_BGR2RGB)
self.writer.add_image(
"{}/canvas_{}.png".format(str(self.env.imgid[i]),
str(step)),
canvas,
log,
)
if step == self.max_episode_length:
for i in range(self.env_batch):
if self.env.imgid[i] < 50:
gt = cv2.cvtColor(
(to_numpy(self.env.gt[i].permute(1, 2, 0))),
cv2.COLOR_BGR2RGB)
canvas = cv2.cvtColor(
(to_numpy(self.env.canvas[i].permute(1, 2, 0))),
cv2.COLOR_BGR2RGB,
)
self.writer.add_image(
str(self.env.imgid[i]) + "/_target.png", gt, log)
self.writer.add_image(
str(self.env.imgid[i]) + "/_canvas.png", canvas, log)
def save_image_with_gen(self, log, step):
for i in range(1):
canvas = to_numpy(self.env.canvas[i].permute(1, 2, 0))
self.writer.add_image(
"images/%d-%d-%d_canvas.png" % (log, i, step), canvas, log)
gt = to_numpy(self.env.gt[i].permute(1, 2, 0))
self.writer.add_image("images/%d-%d_target.png" % (log, i), gt, log)
def select_action(self, state, return_fix=False, noise_factor=0):
self.eval()
with torch.no_grad():
action = self.play(state)
action = to_numpy(action)
if noise_factor > 0:
action = self.noise_action(noise_factor, state, action)
self.train()
self.action = action
if return_fix:
return action
return self.action
def cal_reward(self):
dis = self.cal_dis()
reward = (self.lastdis - dis) / (self.ini_dis + 1e-8)
self.lastdis = dis
return to_numpy(reward)
def get_dist(self):
return to_numpy((((self.env.gt.float() - self.env.canvas.float()) /
255)**2).mean(1).mean(1).mean(1))
def get_grads(self):
"""
Returns the current gradient
"""
return deepcopy(
np.hstack([to_numpy(v.grad).flatten() for v in self.parameters()]))
def get_params(self):
"""
Returns parameters of the actor
"""
return deepcopy(
np.hstack([to_numpy(v).flatten() for v in self.parameters()]))