BerfScene / utils /visualizers.py
3v324v23's picture
init
2f85de4
raw
history blame
25 kB
# python3.7
"""Utility functions for visualizing results."""
import base64
import os.path
import cv2
import numpy as np
from bs4 import BeautifulSoup
__all__ = [
'get_grid_shape', 'get_blank_image', 'load_image', 'save_image',
'resize_image', 'postprocess_image', 'add_text_to_image',
'parse_image_size', 'fuse_images', 'HtmlPageVisualizer', 'HtmlPageReader',
'VideoReader', 'VideoWriter'
]
def get_grid_shape(size, row=0, col=0, is_portrait=False):
"""Gets the shape of a grid based on the size.
This function makes greatest effort on making the output grid square if
neither `row` nor `col` is set. If `is_portrait` is set as `False`, the
height will always be equal to or smaller than the width. For example, if
input `size = 16`, output shape will be `(4, 4)`; if input `size = 15`,
output shape will be (3, 5). Otherwise, the height will always be equal to
or larger than the width.
Args:
size: Size (height * width) of the target grid.
is_portrait: Whether to return a portrait size of a landscape size.
(default: False)
Returns:
A two-element tuple, representing height and width respectively.
"""
assert isinstance(size, int)
assert isinstance(row, int)
assert isinstance(col, int)
if size == 0:
return (0, 0)
if row > 0 and col > 0 and row * col != size:
row = 0
col = 0
if row > 0 and size % row == 0:
return (row, size // row)
if col > 0 and size % col == 0:
return (size // col, col)
row = int(np.sqrt(size))
while row > 0:
if size % row == 0:
col = size // row
break
row = row - 1
return (col, row) if is_portrait else (row, col)
def get_blank_image(height, width, channels=3, is_black=True):
"""Gets a blank image, either white of black.
NOTE: This function will always return an image with `RGB` channel order for
color image and pixel range [0, 255].
Args:
height: Height of the returned image.
width: Width of the returned image.
channels: Number of channels. (default: 3)
is_black: Whether to return a black image. (default: True)
"""
shape = (height, width, channels)
if is_black:
return np.zeros(shape, dtype=np.uint8)
return np.ones(shape, dtype=np.uint8) * 255
def load_image(path, image_channels=3):
"""Loads an image from disk.
NOTE: This function will always return an image with `RGB` channel order for
color image and pixel range [0, 255].
Args:
path: Path to load the image from.
image_channels: Number of image channels of returned image. This field
is employed since `cv2.imread()` will always return a 3-channel
image, even for grayscale image.
Returns:
An image with dtype `np.ndarray`, or `None` if `path` does not exist.
"""
if not os.path.isfile(path):
return None
assert image_channels in [1, 3]
image = cv2.imread(path)
assert image.ndim == 3 and image.shape[2] == 3
if image_channels == 1:
return image[:, :, 0:1]
return image[:, :, ::-1]
def save_image(path, image):
"""Saves an image to disk.
NOTE: The input image (if colorful) is assumed to be with `RGB` channel
order and pixel range [0, 255].
Args:
path: Path to save the image to.
image: Image to save.
"""
if image is None:
return
assert image.ndim == 3 and image.shape[2] in [1, 3]
cv2.imwrite(path, image[:, :, ::-1])
def resize_image(image, *args, **kwargs):
"""Resizes image.
This is a wrap of `cv2.resize()`.
NOTE: THe channel order of the input image will not be changed.
Args:
image: Image to resize.
"""
if image is None:
return None
assert image.ndim == 3 and image.shape[2] in [1, 3]
image = cv2.resize(image, *args, **kwargs)
if image.ndim == 2:
return image[:, :, np.newaxis]
return image
def postprocess_image(image, min_val=-1.0, max_val=1.0, data_format='NCHW'):
"""Post-processes image to pixel range [0, 255] with dtype `uint8`.
NOTE: The returned image will always be with `HWC` format.
Args:
min_val: Minimum value of the input image.
max_val: Maximum value of the input image.
data_format: Data format of the input image. Supporting `NCHW`, `NHWC`,
`CHW`, `HWC`.
Returns:
The post-processed image.
Raises:
NotImplementedError: If the input `data_format` is not support.
"""
assert isinstance(image, np.ndarray)
image = image.astype(np.float64)
image = (image - min_val) * 255 / (max_val - min_val)
image = np.clip(image + 0.5, 0, 255).astype(np.uint8)
data_format = data_format.upper()
if data_format == 'NCHW':
assert image.ndim == 4 and image.shape[1] in [1, 3]
return image.transpose(0, 2, 3, 1)
if data_format == 'NHWC':
assert image.ndim == 4 and image.shape[3] in [1, 3]
return image
if data_format == 'CHW':
assert image.ndim == 3 and image.shape[0] in [1, 3]
return image.transpose(1, 2, 0)
if data_format == 'HWC':
assert image.ndim == 3 and image.shape[2] in [1, 3]
return image
raise NotImplementedError(f'Data format `{data_format}` is not supported!')
def add_text_to_image(image,
text='',
position=None,
font=cv2.FONT_HERSHEY_TRIPLEX,
font_size=1.0,
line_type=cv2.LINE_8,
line_width=1,
color=(255, 255, 255)):
"""Overlays text on given image.
NOTE: The input image is assumed to be with `RGB` channel order.
Args:
image: The image to overlay text on.
text: Text content to overlay on the image. (default: '')
position: Target position (bottom-left corner) to add text. If not set,
center of the image will be used by default. (default: None)
font: Font of the text added. (default: cv2.FONT_HERSHEY_TRIPLEX)
font_size: Font size of the text added. (default: 1.0)
line_type: Line type used to depict the text. (default: cv2.LINE_8)
line_width: Line width used to depict the text. (default: 1)
color: Color of the text added in `RGB` channel order. (default:
(255, 255, 255))
Returns:
An image with target text overlayed on.
"""
if image is None or not text:
return image
cv2.putText(img=image,
text=text,
org=position,
fontFace=font,
fontScale=font_size,
color=color,
thickness=line_width,
lineType=line_type,
bottomLeftOrigin=False)
return image
def parse_image_size(obj):
"""Parses object to a pair of image size, i.e., (width, height).
Args:
obj: The input object to parse image size from.
Returns:
A two-element tuple, indicating image width and height respectively.
Raises:
If the input is invalid, i.e., neither a list or tuple, nor a string.
"""
if obj is None or obj == '':
width = height = 0
elif isinstance(obj, int):
width = height = obj
elif isinstance(obj, (list, tuple, np.ndarray)):
numbers = tuple(obj)
if len(numbers) == 0:
width = height = 0
elif len(numbers) == 1:
width = height = numbers[0]
elif len(numbers) == 2:
width = numbers[0]
height = numbers[1]
else:
raise ValueError(f'At most two elements for image size.')
elif isinstance(obj, str):
splits = obj.replace(' ', '').split(',')
numbers = tuple(map(int, splits))
if len(numbers) == 0:
width = height = 0
elif len(numbers) == 1:
width = height = numbers[0]
elif len(numbers) == 2:
width = numbers[0]
height = numbers[1]
else:
raise ValueError(f'At most two elements for image size.')
else:
raise ValueError(f'Invalid type of input: {type(obj)}!')
return (max(0, width), max(0, height))
def fuse_images(images,
image_size=None,
row=0,
col=0,
is_row_major=True,
is_portrait=False,
row_spacing=0,
col_spacing=0,
border_left=0,
border_right=0,
border_top=0,
border_bottom=0,
black_background=True):
"""Fuses a collection of images into an entire image.
Args:
images: A collection of images to fuse. Should be with shape [num,
height, width, channels].
image_size: This field is used to resize the image before fusion. `0`
disables resizing. (default: None)
row: Number of rows used for image fusion. If not set, this field will
be automatically assigned based on `col` and total number of images.
(default: None)
col: Number of columns used for image fusion. If not set, this field
will be automatically assigned based on `row` and total number of
images. (default: None)
is_row_major: Whether the input images should be arranged row-major or
column-major. (default: True)
is_portrait: Only active when both `row` and `col` should be assigned
automatically. (default: False)
row_spacing: Space between rows. (default: 0)
col_spacing: Space between columns. (default: 0)
border_left: Width of left border. (default: 0)
border_right: Width of right border. (default: 0)
border_top: Width of top border. (default: 0)
border_bottom: Width of bottom border. (default: 0)
Returns:
The fused image.
Raises:
ValueError: If the input `images` is not with shape [num, height, width,
width].
"""
if images is None:
return images
if images.ndim != 4:
raise ValueError(f'Input `images` should be with shape [num, height, '
f'width, channels], but {images.shape} is received!')
num, image_height, image_width, channels = images.shape
width, height = parse_image_size(image_size)
height = height or image_height
width = width or image_width
row, col = get_grid_shape(num, row=row, col=col, is_portrait=is_portrait)
fused_height = (
height * row + row_spacing * (row - 1) + border_top + border_bottom)
fused_width = (
width * col + col_spacing * (col - 1) + border_left + border_right)
fused_image = get_blank_image(
fused_height, fused_width, channels=channels, is_black=black_background)
images = images.reshape(row, col, image_height, image_width, channels)
if not is_row_major:
images = images.transpose(1, 0, 2, 3, 4)
for i in range(row):
y = border_top + i * (height + row_spacing)
for j in range(col):
x = border_left + j * (width + col_spacing)
if height != image_height or width != image_width:
image = cv2.resize(images[i, j], (width, height))
else:
image = images[i, j]
fused_image[y:y + height, x:x + width] = image
return fused_image
def get_sortable_html_header(column_name_list, sort_by_ascending=False):
"""Gets header for sortable html page.
Basically, the html page contains a sortable table, where user can sort the
rows by a particular column by clicking the column head.
Example:
column_name_list = [name_1, name_2, name_3]
header = get_sortable_html_header(column_name_list)
footer = get_sortable_html_footer()
sortable_table = ...
html_page = header + sortable_table + footer
Args:
column_name_list: List of column header names.
sort_by_ascending: Default sorting order. If set as `True`, the html
page will be sorted by ascending order when the header is clicked
for the first time.
Returns:
A string, which represents for the header for a sortable html page.
"""
header = '\n'.join([
'<script type="text/javascript">',
'var column_idx;',
'var sort_by_ascending = ' + str(sort_by_ascending).lower() + ';',
'',
'function sorting(tbody, column_idx){',
' this.column_idx = column_idx;',
' Array.from(tbody.rows)',
' .sort(compareCells)',
' .forEach(function(row) { tbody.appendChild(row); })',
' sort_by_ascending = !sort_by_ascending;',
'}',
'',
'function compareCells(row_a, row_b) {',
' var val_a = row_a.cells[column_idx].innerText;',
' var val_b = row_b.cells[column_idx].innerText;',
' var flag = sort_by_ascending ? 1 : -1;',
' return flag * (val_a > val_b ? 1 : -1);',
'}',
'</script>',
'',
'<html>',
'',
'<head>',
'<style>',
' table {',
' border-spacing: 0;',
' border: 1px solid black;',
' }',
' th {',
' cursor: pointer;',
' }',
' th, td {',
' text-align: left;',
' vertical-align: middle;',
' border-collapse: collapse;',
' border: 0.5px solid black;',
' padding: 8px;',
' }',
' tr:nth-child(even) {',
' background-color: #d2d2d2;',
' }',
'</style>',
'</head>',
'',
'<body>',
'',
'<table>',
'<thead>',
'<tr>',
''])
for idx, name in enumerate(column_name_list):
header += f' <th onclick="sorting(tbody, {idx})">{name}</th>\n'
header += '</tr>\n'
header += '</thead>\n'
header += '<tbody id="tbody">\n'
return header
def get_sortable_html_footer():
"""Gets footer for sortable html page.
Check function `get_sortable_html_header()` for more details.
"""
return '</tbody>\n</table>\n\n</body>\n</html>\n'
def encode_image_to_html_str(image, image_size=None):
"""Encodes an image to html language.
NOTE: Input image is always assumed to be with `RGB` channel order.
Args:
image: The input image to encode. Should be with `RGB` channel order.
image_size: This field is used to resize the image before encoding. `0`
disables resizing. (default: None)
Returns:
A string which represents the encoded image.
"""
if image is None:
return ''
assert image.ndim == 3 and image.shape[2] in [1, 3]
# Change channel order to `BGR`, which is opencv-friendly.
image = image[:, :, ::-1]
# Resize the image if needed.
width, height = parse_image_size(image_size)
if height or width:
height = height or image.shape[0]
width = width or image.shape[1]
image = cv2.resize(image, (width, height))
# Encode the image to html-format string.
encoded_image = cv2.imencode('.jpg', image)[1].tostring()
encoded_image_base64 = base64.b64encode(encoded_image).decode('utf-8')
html_str = f'<img src="data:image/jpeg;base64, {encoded_image_base64}"/>'
return html_str
def decode_html_str_to_image(html_str, image_size=None):
"""Decodes image from html.
Args:
html_str: Image string parsed from html.
image_size: This field is used to resize the image after decoding. `0`
disables resizing. (default: None)
Returns:
An image with `RGB` channel order.
"""
if not html_str:
return None
assert isinstance(html_str, str)
image_str = html_str.split(',')[-1]
encoded_image = base64.b64decode(image_str)
encoded_image_numpy = np.frombuffer(encoded_image, dtype=np.uint8)
image = cv2.imdecode(encoded_image_numpy, flags=cv2.IMREAD_COLOR)
# Resize the image if needed.
width, height = parse_image_size(image_size)
if height or width:
height = height or image.shape[0]
width = width or image.shape[1]
image = cv2.resize(image, (width, height))
return image[:, :, ::-1]
class HtmlPageVisualizer(object):
"""Defines the html page visualizer.
This class can be used to visualize image results as html page. Basically,
it is based on an html-format sorted table with helper functions
`get_sortable_html_header()`, `get_sortable_html_footer()`, and
`encode_image_to_html_str()`. To simplify the usage, specifying the
following fields are enough to create a visualization page:
(1) num_rows: Number of rows of the table (header-row exclusive).
(2) num_cols: Number of columns of the table.
(3) header contents (optional): Title of each column.
NOTE: `grid_size` can be used to assign `num_rows` and `num_cols`
automatically.
Example:
html = HtmlPageVisualizer(num_rows, num_cols)
html.set_headers([...])
for i in range(num_rows):
for j in range(num_cols):
html.set_cell(i, j, text=..., image=..., highlight=False)
html.save('visualize.html')
"""
def __init__(self,
num_rows=0,
num_cols=0,
grid_size=0,
is_portrait=True,
viz_size=None):
if grid_size > 0:
num_rows, num_cols = get_grid_shape(
grid_size, row=num_rows, col=num_cols, is_portrait=is_portrait)
assert num_rows > 0 and num_cols > 0
self.num_rows = num_rows
self.num_cols = num_cols
self.viz_size = parse_image_size(viz_size)
self.headers = ['' for _ in range(self.num_cols)]
self.cells = [[{
'text': '',
'image': '',
'highlight': False,
} for _ in range(self.num_cols)] for _ in range(self.num_rows)]
def set_header(self, col_idx, content):
"""Sets the content of a particular header by column index."""
self.headers[col_idx] = content
def set_headers(self, contents):
"""Sets the contents of all headers."""
if isinstance(contents, str):
contents = [contents]
assert isinstance(contents, (list, tuple))
assert len(contents) == self.num_cols
for col_idx, content in enumerate(contents):
self.set_header(col_idx, content)
def set_cell(self, row_idx, col_idx, text='', image=None, highlight=False):
"""Sets the content of a particular cell.
Basically, a cell contains some text as well as an image. Both text and
image can be empty.
Args:
row_idx: Row index of the cell to edit.
col_idx: Column index of the cell to edit.
text: Text to add into the target cell. (default: None)
image: Image to show in the target cell. Should be with `RGB`
channel order. (default: None)
highlight: Whether to highlight this cell. (default: False)
"""
self.cells[row_idx][col_idx]['text'] = text
self.cells[row_idx][col_idx]['image'] = encode_image_to_html_str(
image, self.viz_size)
self.cells[row_idx][col_idx]['highlight'] = bool(highlight)
def save(self, save_path):
"""Saves the html page."""
html = ''
for i in range(self.num_rows):
html += f'<tr>\n'
for j in range(self.num_cols):
text = self.cells[i][j]['text']
image = self.cells[i][j]['image']
if self.cells[i][j]['highlight']:
color = ' bgcolor="#FF8888"'
else:
color = ''
if text:
html += f' <td{color}>{text}<br><br>{image}</td>\n'
else:
html += f' <td{color}>{image}</td>\n'
html += f'</tr>\n'
header = get_sortable_html_header(self.headers)
footer = get_sortable_html_footer()
with open(save_path, 'w') as f:
f.write(header + html + footer)
class HtmlPageReader(object):
"""Defines the html page reader.
This class can be used to parse results from the visualization page
generated by `HtmlPageVisualizer`.
Example:
html = HtmlPageReader(html_path)
for j in range(html.num_cols):
header = html.get_header(j)
for i in range(html.num_rows):
for j in range(html.num_cols):
text = html.get_text(i, j)
image = html.get_image(i, j, image_size=None)
"""
def __init__(self, html_path):
"""Initializes by loading the content from file."""
self.html_path = html_path
if not os.path.isfile(html_path):
raise ValueError(f'File `{html_path}` does not exist!')
# Load content.
with open(html_path, 'r') as f:
self.html = BeautifulSoup(f, 'html.parser')
# Parse headers.
thead = self.html.find('thead')
headers = thead.findAll('th')
self.headers = []
for header in headers:
self.headers.append(header.text)
self.num_cols = len(self.headers)
# Parse cells.
tbody = self.html.find('tbody')
rows = tbody.findAll('tr')
self.cells = []
for row in rows:
cells = row.findAll('td')
self.cells.append([])
for cell in cells:
self.cells[-1].append({
'text': cell.text,
'image': cell.find('img')['src'],
})
assert len(self.cells[-1]) == self.num_cols
self.num_rows = len(self.cells)
def get_header(self, j):
"""Gets header for a particular column."""
return self.headers[j]
def get_text(self, i, j):
"""Gets text from a particular cell."""
return self.cells[i][j]['text']
def get_image(self, i, j, image_size=None):
"""Gets image from a particular cell."""
return decode_html_str_to_image(self.cells[i][j]['image'], image_size)
class VideoReader(object):
"""Defines the video reader.
This class can be used to read frames from a given video.
"""
def __init__(self, path):
"""Initializes the video reader by loading the video from disk."""
if not os.path.isfile(path):
raise ValueError(f'Video `{path}` does not exist!')
self.path = path
self.video = cv2.VideoCapture(path)
assert self.video.isOpened()
self.position = 0
self.length = int(self.video.get(cv2.CAP_PROP_FRAME_COUNT))
self.frame_height = int(self.video.get(cv2.CAP_PROP_FRAME_HEIGHT))
self.frame_width = int(self.video.get(cv2.CAP_PROP_FRAME_WIDTH))
self.fps = self.video.get(cv2.CAP_PROP_FPS)
def __del__(self):
"""Releases the opened video."""
self.video.release()
def read(self, position=None):
"""Reads a certain frame.
NOTE: The returned frame is assumed to be with `RGB` channel order.
Args:
position: Optional. If set, the reader will read frames from the
exact position. Otherwise, the reader will read next frames.
(default: None)
"""
if position is not None and position < self.length:
self.video.set(cv2.CAP_PROP_POS_FRAMES, position)
self.position = position
success, frame = self.video.read()
self.position = self.position + 1
return frame[:, :, ::-1] if success else None
class VideoWriter(object):
"""Defines the video writer.
This class can be used to create a video.
NOTE: `.avi` and `DIVX` is the most recommended codec format since it does
not rely on other dependencies.
"""
def __init__(self, path, frame_height, frame_width, fps=24, codec='DIVX'):
"""Creates the video writer."""
self.path = path
self.frame_height = frame_height
self.frame_width = frame_width
self.fps = fps
self.codec = codec
self.video = cv2.VideoWriter(filename=path,
fourcc=cv2.VideoWriter_fourcc(*codec),
fps=fps,
frameSize=(frame_width, frame_height))
def __del__(self):
"""Releases the opened video."""
self.video.release()
def write(self, frame):
"""Writes a target frame.
NOTE: The input frame is assumed to be with `RGB` channel order.
"""
self.video.write(frame[:, :, ::-1])