# Path: util/util_visualize.py
# def convert2rgb(tensor, key, **kargs):
#     found = False
#     for k in KEYWORDS:
#         if k in key:
#             convert = getattr(Converter, k + '2img')
#             found = True
#             break
#     if not found:
#         return None
#     else:
#         return convert(tensor, **kargs)
# 
# def is_key_image(key):
#     """check if the given key correspondes to images
# 
#     Arguments:
#         key {str} -- key of a data pack
# 
#     Returns:
#         bool -- [True if the given key correspondes to an image]
#     """
# 
#     for k in KEYWORDS:
#         if k in key:
#             return True
#     return False
# 
# Path: util/util_flow.py
# def flow2img(flow_data, maxrad=None):
#     """
#     convert optical flow into color image
#     :param flow_data:
#     :return: color image
#     """
#     # print(flow_data.shape)
#     # print(type(flow_data))
#     u = flow_data[:, :, 0]
#     v = flow_data[:, :, 1]
# 
#     UNKNOW_FLOW_THRESHOLD = 1e7
#     pr1 = abs(u) > UNKNOW_FLOW_THRESHOLD
#     pr2 = abs(v) > UNKNOW_FLOW_THRESHOLD
#     idx_unknown = (pr1 | pr2)
#     u[idx_unknown] = v[idx_unknown] = 0
# 
#     # get max value in each direction
#     '''
#     maxu = -999.
#     maxv = -999.
#     minu = 999.
#     minv = 999.
#     maxu = max(maxu, np.max(u))
#     maxv = max(maxv, np.max(v))
#     minu = min(minu, np.min(u))
#     minv = min(minv, np.min(v))
#     '''
# 
#     rad = np.sqrt(u ** 2 + v ** 2)
#     if maxrad is None:
#         maxrad = max(-1, np.max(rad))
#     u = u / maxrad + np.finfo(float).eps
#     v = v / maxrad + np.finfo(float).eps
# 
#     img = compute_color(u, v)
# 
#     idx = np.repeat(idx_unknown[:, :, np.newaxis], 3, axis=2)
#     img[idx] = 0
# 
#     return img


from util.util_visualize import convert2rgb, is_key_image
from torch.multiprocessing import Pool
from util.util_flow import flow2img
from os.path import join
from PIL import Image
import atexit
import numpy as np
                pool.close()
                pool.join()
        atexit.register(cleanup)

    def visualize(self, pack, batch_idx, outdir, is_test=False):
        # first append to the shared content
        # then launch the subprocesses for dumping images.
        # b_size = pack['batch_size']
        epoch_folder = outdir.split('/')[-1]
        if is_test:
            epoch_folder = None
        self.prepare_HTML_string(pack, batch_idx, epoch_folder)

        if self.pool:
            self.pool.apply_async(
                self._visualize,
                [pack, batch_idx, outdir],
                error_callback=self._error_callback
            )
        else:
            self._visualize(pack, batch_idx, outdir)
        # prepare HTML string

    def prepare_HTML_string(self, pack, batch_idx, epoch_folder):
        if self.html_logger.epoch_content is None:
            self.html_logger.epoch_content = {}
            #self.html_logger.epoch_content['header'] = ''
            header = ''
            for k in sorted(list(pack.keys())):
                # get the ones that are useful:

# Path: util/util_visualize.py
# def convert2rgb(tensor, key, **kargs):
#     found = False
#     for k in KEYWORDS:
#         if k in key:
#             convert = getattr(Converter, k + '2img')
#             found = True
#             break
#     if not found:
#         return None
#     else:
#         return convert(tensor, **kargs)
# 
# def is_key_image(key):
#     """check if the given key correspondes to images
# 
#     Arguments:
#         key {str} -- key of a data pack
# 
#     Returns:
#         bool -- [True if the given key correspondes to an image]
#     """
# 
#     for k in KEYWORDS:
#         if k in key:
#             return True
#     return False
# 
# Path: util/util_flow.py
# def flow2img(flow_data, maxrad=None):
#     """
#     convert optical flow into color image
#     :param flow_data:
#     :return: color image
#     """
#     # print(flow_data.shape)
#     # print(type(flow_data))
#     u = flow_data[:, :, 0]
#     v = flow_data[:, :, 1]
# 
#     UNKNOW_FLOW_THRESHOLD = 1e7
#     pr1 = abs(u) > UNKNOW_FLOW_THRESHOLD
#     pr2 = abs(v) > UNKNOW_FLOW_THRESHOLD
#     idx_unknown = (pr1 | pr2)
#     u[idx_unknown] = v[idx_unknown] = 0
# 
#     # get max value in each direction
#     '''
#     maxu = -999.
#     maxv = -999.
#     minu = 999.
#     minv = 999.
#     maxu = max(maxu, np.max(u))
#     maxv = max(maxv, np.max(v))
#     minu = min(minu, np.min(u))
#     minv = min(minv, np.min(v))
#     '''
# 
#     rad = np.sqrt(u ** 2 + v ** 2)
#     if maxrad is None:
#         maxrad = max(-1, np.max(rad))
#     u = u / maxrad + np.finfo(float).eps
#     v = v / maxrad + np.finfo(float).eps
# 
#     img = compute_color(u, v)
# 
#     idx = np.repeat(idx_unknown[:, :, np.newaxis], 3, axis=2)
#     img[idx] = 0
# 
#     return img


# Path: visualize/html_visualizer.py
from util.util_visualize import convert2rgb, is_key_image
from torch.multiprocessing import Pool
from util.util_flow import flow2img
from os.path import join
from PIL import Image
import atexit
import numpy as np

                pool.close()
                pool.join()
        atexit.register(cleanup)

    def visualize(self, pack, batch_idx, outdir, is_test=False):
        # first append to the shared content
        # then launch the subprocesses for dumping images.
        # b_size = pack['batch_size']
        epoch_folder = outdir.split('/')[-1]
        if is_test:
            epoch_folder = None
        self.prepare_HTML_string(pack, batch_idx, epoch_folder)

        if self.pool:
            self.pool.apply_async(
                self._visualize,
                [pack, batch_idx, outdir],
                error_callback=self._error_callback
            )
        else:
            self._visualize(pack, batch_idx, outdir)
        # prepare HTML string

    def prepare_HTML_string(self, pack, batch_idx, epoch_folder):
        if self.html_logger.epoch_content is None:
            self.html_logger.epoch_content = {}
            #self.html_logger.epoch_content['header'] = ''
            header = ''
            for k in sorted(list(pack.keys())):
                # get the ones that are useful:

                if is_key_image(k):