<reponame>ojmakhura/DIGITS<filename>digits/extensions/view/imageOutput/view.py
# Copyright (c) 2016-2017, NVIDIA CORPORATION.  All rights reserved.


import os

import PIL.Image
import PIL.ImageDraw

import digits
from digits.utils import subclass, override
from .forms import ConfigForm
from ..interface import VisualizationInterface

CONFIG_TEMPLATE = "config_template.html"
VIEW_TEMPLATE = "view_template.html"


@subclass
class Visualization(VisualizationInterface):
    """
    A visualization extension to display the network output as an image
    """

    def __init__(self, dataset, **kwargs):
        # memorize view template for later use
        extension_dir = os.path.dirname(os.path.abspath(__file__))
        self.view_template = open(
            os.path.join(extension_dir, VIEW_TEMPLATE), "r").read()

        # view options
        self.channel_order = kwargs['channel_order'].upper()
        self.data_order = kwargs['data_order'].upper()
        self.normalize = (kwargs['pixel_conversion'] == 'normalize')
        self.show_input = (kwargs['show_input'] == 'yes')

    @staticmethod
    def get_config_form():
        return ConfigForm()

    @staticmethod
    def get_config_template(form):
        """
        parameters:
        - form: form returned by get_config_form(). This may be populated
        with values if the job was cloned
        returns:
        - (template, context) tuple
          - template is a Jinja template to use for rendering config options
          - context is a dictionary of context variables to use for rendering
          the form
        """
        extension_dir = os.path.dirname(os.path.abspath(__file__))
        template = open(
            os.path.join(extension_dir, CONFIG_TEMPLATE), "r").read()
        return (template, {'form': form})

    @staticmethod
    def get_id():
        return 'image-image-output'

    @staticmethod
    def get_title():
        return 'Image output'

    @override
    def get_view_template(self, data):
        """
        returns:
        - (template, context) tuple
          - template is a Jinja template to use for rendering config options
          - context is a dictionary of context variables to use for rendering
          the form
        """
        return self.view_template, {'image_input': digits.utils.image.embed_image_html(data[0]),
                                    'image_output': digits.utils.image.embed_image_html(data[1])}

    @override
    def process_data(self, input_id, input_data, output_data):
        """
        Process one inference and return data to visualize
        """

        if self.show_input:
            data_input = input_data.astype('float32')
            image_input = self.process_image(self.data_order, data_input)
        else:
            image_input = None

        data_output = output_data[list(output_data.keys())[0]].astype('float32')
        image_output = self.process_image(self.data_order, data_output)

        return [image_input, image_output]

    def process_image(self, data_order, data):
        if data_order == 'HWC':
            data = (data.transpose((2, 0, 1)))

        # assume CHW at this point
        channels = data.shape[0]
        if channels == 3 and self.channel_order == 'BGR':
            data = data[[2, 1, 0], ...]  # BGR to RGB
        # convert to HWC
        data = data.transpose((1, 2, 0))
        # assume 8-bit
        if self.normalize:
            data -= data.min()
            if data.max() > 0:
                data /= data.max()
                data *= 255
        else:
            # clip
            data = data.clip(0, 255)
        # convert to uint8
        data = data.astype('uint8')
        # convert to PIL image
        if channels == 1:
            # drop channel axis
            image = PIL.Image.fromarray(data[:, :, 0])
        elif channels == 3:
            image = PIL.Image.fromarray(data)
        else:
            raise ValueError("Unhandled number of channels: %d" % channels)

        return image
