Spaces:
Sleeping
Sleeping
File size: 3,445 Bytes
c64c726 |
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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
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)
|