Spaces:
Sleeping
Sleeping
| from PIL import Image, ImageDraw, ImageFont | |
| import numpy as np | |
| class FrameRenderer: | |
| """ | |
| Helper class for drawing headers, text, and observation images onto a PIL.Image frame. | |
| """ | |
| def __init__( | |
| self, | |
| frame_width, | |
| frame_height, | |
| header_rect, | |
| obs_rect, | |
| obs_low_res_rect=None, | |
| font_size=16, | |
| font_path=None, | |
| ): | |
| """ | |
| frame_width, frame_height: output frame size | |
| header_rect: (x, y, w, h) region for header | |
| obs_rect: (x, y, w, h) region for main observation | |
| obs_low_res_rect: (x, y, w, h) region for low-res obs (optional) | |
| font_size: font size for text | |
| font_path: path to .ttf font file (optional) | |
| """ | |
| self.frame_width = frame_width | |
| self.frame_height = frame_height | |
| self.header_rect = header_rect | |
| self.obs_rect = obs_rect | |
| self.obs_low_res_rect = obs_low_res_rect | |
| self.font_size = font_size | |
| if font_path: | |
| self.font = ImageFont.truetype(font_path, font_size) | |
| else: | |
| self.font = ImageFont.load_default() | |
| def clear_header(self, image): | |
| """ | |
| Draw a black rectangle with white outline at the header area of image. | |
| """ | |
| draw = ImageDraw.Draw(image) | |
| x, y, w, h = self.header_rect | |
| draw.rectangle([x, y, x + w, y + h], fill="black", outline="white", width=1) | |
| def draw_text(self, image, text, idx_line, idx_column, num_cols): | |
| """ | |
| Draw text at (row, column) in the header area of image. | |
| """ | |
| x_header, y_header, header_width, header_height = self.header_rect | |
| x_pos = 5 + idx_column * int(header_width // num_cols) | |
| y_pos = 5 + idx_line * self.font_size | |
| assert (0 <= x_pos <= header_width) and (0 <= y_pos <= header_height) | |
| draw = ImageDraw.Draw(image) | |
| draw.text((x_header + x_pos, y_header + y_pos), text, font=self.font, fill="white") | |
| def draw_obs(self, image, obs, obs_low_res=None): | |
| """ | |
| Draw the main observation (and optional low-res obs) onto the image. | |
| """ | |
| # Main observation | |
| assert obs.ndim == 4 and obs.size(0) == 1 | |
| obs_img = Image.fromarray( | |
| obs[0].add(1).div(2).mul(255).byte().permute(1, 2, 0).cpu().numpy() | |
| ) | |
| x, y, w, h = self.obs_rect | |
| obs_img = obs_img.resize((w, h), resample=Image.BICUBIC) | |
| image.paste(obs_img, (x, y)) | |
| # Small resolution observation (if any) | |
| if obs_low_res is not None and self.obs_low_res_rect is not None: | |
| assert obs_low_res.ndim == 4 and obs_low_res.size(0) == 1 | |
| obs_lr_img = Image.fromarray( | |
| obs_low_res[0].add(1).div(2).mul(255).byte().permute(1, 2, 0).cpu().numpy() | |
| ) | |
| x_lr, y_lr, w_lr, h_lr = self.obs_low_res_rect | |
| obs_lr_img = obs_lr_img.resize((w_lr, h_lr), resample=Image.BICUBIC) | |
| image.paste(obs_lr_img, (x_lr, y_lr)) | |
| def draw_header_and_text(self, image, header): | |
| """ | |
| Draw header grid and all text in the header area. | |
| header: list of columns, each column is a list of text (rows) | |
| """ | |
| self.clear_header(image) | |
| num_cols = len(header) | |
| for j, col in enumerate(header): | |
| for i, row in enumerate(col): | |
| self.draw_text(image, row, idx_line=i, idx_column=j, num_cols=num_cols) | |