HTML table for wandb that supports images

wandb_htmltable.py

'''
Workaround for logging a simple table that supports step sliding. (See issue https://github.com/wandb/wandb/issues/6286)
It's a great pity that wandb currently doesn't support this with the `wandb.Table` which is too overkill.

The `wandb_htmltable` function follows the same signature as `wandb.Table` and takes as input parameters of the same type.
It currently only supports text and image type data. Image data is realized via its byte string declared in the <img /> tag

Example:
```
  my_data = [
      [0, PIL_img_1, 0, 0],
      [1, PIL_img_2, 8, 0],
      [2, PIL_img_3, 7, 1],
      [3, PIL_img_4, 1, 1],
  ]
  columns = ["id", "image", "prediction", "truth"]
  test_table = wandb_htmltable(data=my_data, columns=columns)
  wandb.log({'test_table': test_table})
```

In addition, controls have been added to specify CSS. The `css` argument for `wandb_htmltable` will end up in the <style> tag of the iframe html.
Each individual data can also be specified with their own inline styles using a dictionary with keys 'content' and 'style'.
The value for 'content' is the data itself and the value for 'style' is the inline style in string type.

Example:
```
  my_data = [
      [{'content': 0, 'style': 'font-weight: bold;'}, {'content': PIL_img_1, 'style': 'max-height:512px; max-width:512px'}, 0, 0],
      [{'content': 1, 'style': 'font-weight: bold;'}, {'content': PIL_img_2, 'style': 'max-height:512px; max-width:512px'}, 8, 0],
      [{'content': 2, 'style': 'font-weight: bold;'}, {'content': PIL_img_3, 'style': 'max-height:512px; max-width:512px'}, 7, 1],
      [{'content': 3, 'style': 'font-weight: bold;'}, {'content': PIL_img_4, 'style': 'max-height:512px; max-width:512px'}, 1, 1],
  ]
```
'''

import io
import wandb
import base64
from PIL import Image

def wandb_htmltable(data, columns, css=''):
	# Helper functions
	def check_shape():
		num_cols = len(columns)
		for i,row in enumerate(data):
			if len(row) != num_cols:
				raise ValueError(f'Row {i} in data should have {num_cols} cols, but found to have {len(row)}.')
	def join(list_str):
		return "".join(list_str)
	def wrap_html_tag(tag, content):
		html_content = content
		inline_style = ''
		if type(content) == dict:
			html_content = content['content']
			inline_style = content.get('style', inline_style)
		if Image.isImageType(html_content):	# if PIL image
			html_content = format_image(html_content, inline_style)
		return f'<{tag} style="{inline_style}">{html_content}</{tag}>'
	def format_image(pil_image, inline_style):
		img_byte_arr = io.BytesIO()
		pil_image.save(img_byte_arr, format='PNG')
		btyestr = base64.b64encode(img_byte_arr.getvalue()).decode()
		return f'<img style="{inline_style}" src="data:image/png;base64,{btyestr}" />'
	def format_html_tr(list_data, is_header=False):
		return wrap_html_tag('tr', join([format_html_td(cd, is_header) for cd in list_data]))
	def format_html_td(cell_data, is_header):
		return wrap_html_tag('th' if is_header else 'td', cell_data)

	check_shape()
	if not css:
		css = wrap_html_tag('style', '''
		table, th, td {
			border: 1px solid black;
			border-collapse: collapse;
		}''')
	column_html = format_html_tr(columns, is_header=True)
	rows_html = join([format_html_tr(row) for row in data])
	table_html = wrap_html_tag('table', join([column_html, rows_html]))
	body_html = wrap_html_tag('body', table_html)
	head_html = wrap_html_tag('head', css)
	iframe_html = wrap_html_tag('html', join([head_html, body_html]))
	return wandb.Html(iframe_html, inject=True)