#!/usr/bin/env python3
"""
 Copyright (c) 2018-2021 Intel Corporation

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
"""

from openvino.runtime import Core, get_version
import cv2 as cv
import numpy as np
import logging as log
from time import perf_counter
import sys
from argparse import ArgumentParser, SUPPRESS
from pathlib import Path

sys.path.append(str(Path(__file__).resolve().parents[2] / 'common/python'))
sys.path.append(str(Path(__file__).resolve().parents[2] / 'common/python/openvino/model_zoo'))

import monitors
from images_capture import open_images_capture
from model_api.performance_metrics import PerformanceMetrics

log.basicConfig(format='[ %(levelname)s ] %(message)s', level=log.DEBUG, stream=sys.stdout)


def build_arg():
    parser = ArgumentParser(add_help=False)
    in_args = parser.add_argument_group('Options')
    in_args.add_argument('-h', '--help', action='help', default=SUPPRESS, help='Help with the script.')
    in_args.add_argument("-m", "--model", help="Required. Path to .xml file with pre-trained model.",
                         required=True, type=Path)
    in_args.add_argument("-d", "--device",
                         help="Optional. Specify target device for infer: CPU, GPU, HDDL or MYRIAD. "
                              "Default: CPU",
                         default="CPU", type=str)
    in_args.add_argument('-i', "--input", required=True,
                         help='Required. An input to process. The input must be a single image, '
                              'a folder of images, video file or camera id.')
    in_args.add_argument('--loop', default=False, action='store_true',
                         help='Optional. Enable reading the input in a loop.')
    in_args.add_argument('-o', '--output', required=False,
                         help='Optional. Name of the output file(s) to save.')
    in_args.add_argument('-limit', '--output_limit', required=False, default=1000, type=int,
                         help='Optional. Number of frames to store in output. '
                              'If 0 is set, all frames are stored.')
    in_args.add_argument("--no_show", help="Optional. Don't show output.",
                         action='store_true', default=False)
    in_args.add_argument("-u", "--utilization_monitors", default="", type=str,
                         help="Optional. List of monitors to show initially.")
    return parser

def main(args):
    cap = open_images_capture(args.input, args.loop)

    log.info('OpenVINO Inference Engine')
    log.info('\tbuild: {}'.format(get_version()))
    core = Core()

    log.info('Reading model {}'.format(args.model))
    model = core.read_model(args.model, args.model.with_suffix(".bin"))

    input_tensor_name = 'data_l'
    input_shape = model.input(input_tensor_name).shape
    assert input_shape[1] == 1, "Expected model input shape with 1 channel"

    inputs = {}
    for input in model.inputs:
        inputs[input.get_any_name()] = np.zeros(input.shape)

    assert len(model.outputs) == 1, "Expected number of outputs is equal 1"

    compiled_model = core.compile_model(model, device_name=args.device)
    infer_request = compiled_model.create_infer_request()
    log.info('The model {} is loaded to {}'.format(args.model, args.device))

    _, _, h_in, w_in = input_shape

    frames_processed = 0
    imshow_size = (640, 480)
    graph_size = (imshow_size[0] // 2, imshow_size[1] // 4)
    presenter = monitors.Presenter(args.utilization_monitors, imshow_size[1] * 2 - graph_size[1], graph_size)
    metrics = PerformanceMetrics()

    video_writer = cv.VideoWriter()
    if args.output and not video_writer.open(args.output, cv.VideoWriter_fourcc(*'MJPG'),
                                             cap.fps(), (imshow_size[0] * 2, imshow_size[1] * 2)):
        raise RuntimeError("Can't open video writer")

    start_time = perf_counter()
    original_frame = cap.read()
    if original_frame is None:
        raise RuntimeError("Can't read an image from the input")

    while original_frame is not None:
        (h_orig, w_orig) = original_frame.shape[:2]

        if original_frame.shape[2] > 1:
            frame = cv.cvtColor(cv.cvtColor(original_frame, cv.COLOR_BGR2GRAY), cv.COLOR_GRAY2RGB)
        else:
            frame = cv.cvtColor(original_frame, cv.COLOR_GRAY2RGB)

        img_rgb = frame.astype(np.float32) / 255
        img_lab = cv.cvtColor(img_rgb, cv.COLOR_RGB2Lab)
        img_l_rs = cv.resize(img_lab.copy(), (w_in, h_in))[:, :, 0]

        inputs[input_tensor_name] = np.expand_dims(img_l_rs, axis=[0, 1])

        res = next(iter(infer_request.infer(inputs).values()))

        update_res = np.squeeze(res)

        out = update_res.transpose((1, 2, 0))
        out = cv.resize(out, (w_orig, h_orig))
        img_lab_out = np.concatenate((img_lab[:, :, 0][:, :, np.newaxis], out), axis=2)
        img_bgr_out = np.clip(cv.cvtColor(img_lab_out, cv.COLOR_Lab2BGR), 0, 1)

        original_image = cv.resize(original_frame, imshow_size)
        grayscale_image = cv.resize(frame, imshow_size)
        colorize_image = (cv.resize(img_bgr_out, imshow_size) * 255).astype(np.uint8)
        lab_image = cv.resize(img_lab_out, imshow_size).astype(np.uint8)

        original_image = cv.putText(original_image, 'Original', (25, 50),
                                    cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv.LINE_AA)
        grayscale_image = cv.putText(grayscale_image, 'Grayscale', (25, 50),
                                     cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv.LINE_AA)
        colorize_image = cv.putText(colorize_image, 'Colorize', (25, 50),
                                    cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv.LINE_AA)
        lab_image = cv.putText(lab_image, 'LAB interpretation', (25, 50),
                               cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv.LINE_AA)

        ir_image = [cv.hconcat([original_image, grayscale_image]),
                    cv.hconcat([lab_image, colorize_image])]
        final_image = cv.vconcat(ir_image)

        metrics.update(start_time, final_image)

        frames_processed += 1
        if video_writer.isOpened() and (args.output_limit <= 0 or frames_processed <= args.output_limit):
            video_writer.write(final_image)

        presenter.drawGraphs(final_image)
        if not args.no_show:
            cv.imshow('Colorization Demo', final_image)
            key = cv.waitKey(1)
            if key in {ord("q"), ord("Q"), 27}:
                break
            presenter.handleKey(key)
        start_time = perf_counter()
        original_frame = cap.read()

    metrics.log_total()
    for rep in presenter.reportMeans():
        log.info(rep)

if __name__ == "__main__":
    args = build_arg().parse_args()
    sys.exit(main(args) or 0)

[((34, 0, 34, 91), 'logging.basicConfig', 'log.basicConfig', (), '', True, 'import logging as log\n'), ((38, 13, 38, 43), 'argparse.ArgumentParser', 'ArgumentParser', (), '', False, 'from argparse import ArgumentParser, SUPPRESS\n'), ((64, 10, 64, 52), 'images_capture.open_images_capture', 'open_images_capture', ({(64, 30, 64, 40): 'args.input', (64, 42, 64, 51): 'args.loop'}, {}), '(args.input, args.loop)', False, 'from images_capture import open_images_capture\n'), ((66, 4, 66, 41), 'logging.info', 'log.info', ({(66, 13, 66, 40): '"""OpenVINO Inference Engine"""'}, {}), "('OpenVINO Inference Engine')", True, 'import logging as log\n'), ((68, 11, 68, 17), 'openvino.runtime.Core', 'Core', ({}, {}), '()', False, 'from openvino.runtime import Core, get_version\n'), ((92, 16, 92, 109), 'monitors.Presenter', 'monitors.Presenter', ({(92, 35, 92, 60): 'args.utilization_monitors', (92, 62, 92, 96): 'imshow_size[1] * 2 - graph_size[1]', (92, 98, 92, 108): 'graph_size'}, {}), '(args.utilization_monitors, imshow_size[1] * 2 -\n    graph_size[1], graph_size)', False, 'import monitors\n'), ((93, 14, 93, 34), 'model_api.performance_metrics.PerformanceMetrics', 'PerformanceMetrics', ({}, {}), '()', False, 'from model_api.performance_metrics import PerformanceMetrics\n'), ((95, 19, 95, 35), 'cv2.VideoWriter', 'cv.VideoWriter', ({}, {}), '()', True, 'import cv2 as cv\n'), ((100, 17, 100, 31), 'time.perf_counter', 'perf_counter', ({}, {}), '()', False, 'from time import perf_counter\n'), ((79, 39, 79, 60), 'numpy.zeros', 'np.zeros', ({(79, 48, 79, 59): 'input.shape'}, {}), '(input.shape)', True, 'import numpy as np\n'), ((114, 18, 114, 56), 'cv2.cvtColor', 'cv.cvtColor', ({(114, 30, 114, 37): 'img_rgb', (114, 39, 114, 55): 'cv.COLOR_RGB2Lab'}, {}), '(img_rgb, cv.COLOR_RGB2Lab)', True, 'import cv2 as cv\n'), ((117, 36, 117, 73), 'numpy.expand_dims', 'np.expand_dims', (), '', True, 'import numpy as np\n'), ((121, 21, 121, 36), 'numpy.squeeze', 'np.squeeze', ({(121, 32, 121, 35): 'res'}, {}), '(res)', True, 'import numpy as np\n'), ((124, 14, 124, 46), 'cv2.resize', 'cv.resize', ({(124, 24, 124, 27): 'out', (124, 29, 124, 45): '(w_orig, h_orig)'}, {}), '(out, (w_orig, h_orig))', True, 'import cv2 as cv\n'), ((125, 22, 125, 87), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((128, 25, 128, 63), 'cv2.resize', 'cv.resize', ({(128, 35, 128, 49): 'original_frame', (128, 51, 128, 62): 'imshow_size'}, {}), '(original_frame, imshow_size)', True, 'import cv2 as cv\n'), ((129, 26, 129, 55), 'cv2.resize', 'cv.resize', ({(129, 36, 129, 41): 'frame', (129, 43, 129, 54): 'imshow_size'}, {}), '(frame, imshow_size)', True, 'import cv2 as cv\n'), ((133, 25, 134, 91), 'cv2.putText', 'cv.putText', ({(133, 36, 133, 50): 'original_image', (133, 52, 133, 62): '"""Original"""', (133, 64, 133, 72): '(25, 50)', (134, 36, 134, 59): 'cv.FONT_HERSHEY_SIMPLEX', (134, 61, 134, 62): '1', (134, 64, 134, 75): '(0, 0, 255)', (134, 77, 134, 78): '2', (134, 80, 134, 90): 'cv.LINE_AA'}, {}), "(original_image, 'Original', (25, 50), cv.FONT_HERSHEY_SIMPLEX, 1,\n    (0, 0, 255), 2, cv.LINE_AA)", True, 'import cv2 as cv\n'), ((135, 26, 136, 92), 'cv2.putText', 'cv.putText', ({(135, 37, 135, 52): 'grayscale_image', (135, 54, 135, 65): '"""Grayscale"""', (135, 67, 135, 75): '(25, 50)', (136, 37, 136, 60): 'cv.FONT_HERSHEY_SIMPLEX', (136, 62, 136, 63): '1', (136, 65, 136, 76): '(0, 0, 255)', (136, 78, 136, 79): '2', (136, 81, 136, 91): 'cv.LINE_AA'}, {}), "(grayscale_image, 'Grayscale', (25, 50), cv.FONT_HERSHEY_SIMPLEX,\n    1, (0, 0, 255), 2, cv.LINE_AA)", True, 'import cv2 as cv\n'), ((137, 25, 138, 91), 'cv2.putText', 'cv.putText', ({(137, 36, 137, 50): 'colorize_image', (137, 52, 137, 62): '"""Colorize"""', (137, 64, 137, 72): '(25, 50)', (138, 36, 138, 59): 'cv.FONT_HERSHEY_SIMPLEX', (138, 61, 138, 62): '1', (138, 64, 138, 75): '(0, 0, 255)', (138, 77, 138, 78): '2', (138, 80, 138, 90): 'cv.LINE_AA'}, {}), "(colorize_image, 'Colorize', (25, 50), cv.FONT_HERSHEY_SIMPLEX, 1,\n    (0, 0, 255), 2, cv.LINE_AA)", True, 'import cv2 as cv\n'), ((139, 20, 140, 86), 'cv2.putText', 'cv.putText', ({(139, 31, 139, 40): 'lab_image', (139, 42, 139, 62): '"""LAB interpretation"""', (139, 64, 139, 72): '(25, 50)', (140, 31, 140, 54): 'cv.FONT_HERSHEY_SIMPLEX', (140, 56, 140, 57): '1', (140, 59, 140, 70): '(0, 0, 255)', (140, 72, 140, 73): '2', (140, 75, 140, 85): 'cv.LINE_AA'}, {}), "(lab_image, 'LAB interpretation', (25, 50), cv.\n    FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv.LINE_AA)", True, 'import cv2 as cv\n'), ((144, 22, 144, 42), 'cv2.vconcat', 'cv.vconcat', ({(144, 33, 144, 41): 'ir_image'}, {}), '(ir_image)', True, 'import cv2 as cv\n'), ((159, 21, 159, 35), 'time.perf_counter', 'perf_counter', ({}, {}), '()', False, 'from time import perf_counter\n'), ((164, 8, 164, 21), 'logging.info', 'log.info', ({(164, 17, 164, 20): 'rep'}, {}), '(rep)', True, 'import logging as log\n'), ((67, 34, 67, 47), 'openvino.runtime.get_version', 'get_version', ({}, {}), '()', False, 'from openvino.runtime import Core, get_version\n'), ((111, 20, 111, 66), 'cv2.cvtColor', 'cv.cvtColor', ({(111, 32, 111, 46): 'original_frame', (111, 48, 111, 65): 'cv.COLOR_GRAY2RGB'}, {}), '(original_frame, cv.COLOR_GRAY2RGB)', True, 'import cv2 as cv\n'), ((126, 30, 126, 72), 'cv2.cvtColor', 'cv.cvtColor', ({(126, 42, 126, 53): 'img_lab_out', (126, 55, 126, 71): 'cv.COLOR_Lab2BGR'}, {}), '(img_lab_out, cv.COLOR_Lab2BGR)', True, 'import cv2 as cv\n'), ((142, 20, 142, 65), 'cv2.hconcat', 'cv.hconcat', ({(142, 31, 142, 64): '[original_image, grayscale_image]'}, {}), '([original_image, grayscale_image])', True, 'import cv2 as cv\n'), ((143, 20, 143, 59), 'cv2.hconcat', 'cv.hconcat', ({(143, 31, 143, 58): '[lab_image, colorize_image]'}, {}), '([lab_image, colorize_image])', True, 'import cv2 as cv\n'), ((154, 12, 154, 55), 'cv2.imshow', 'cv.imshow', ({(154, 22, 154, 41): '"""Colorization Demo"""', (154, 43, 154, 54): 'final_image'}, {}), "('Colorization Demo', final_image)", True, 'import cv2 as cv\n'), ((155, 18, 155, 31), 'cv2.waitKey', 'cv.waitKey', ({(155, 29, 155, 30): '1'}, {}), '(1)', True, 'import cv2 as cv\n'), ((96, 58, 96, 88), 'cv2.VideoWriter_fourcc', 'cv.VideoWriter_fourcc', ({(96, 80, 96, 87): "*'MJPG'"}, {}), "(*'MJPG')", True, 'import cv2 as cv\n'), ((109, 32, 109, 78), 'cv2.cvtColor', 'cv.cvtColor', ({(109, 44, 109, 58): 'original_frame', (109, 60, 109, 77): 'cv.COLOR_BGR2GRAY'}, {}), '(original_frame, cv.COLOR_BGR2GRAY)', True, 'import cv2 as cv\n'), ((131, 20, 131, 55), 'cv2.resize', 'cv.resize', ({(131, 30, 131, 41): 'img_lab_out', (131, 43, 131, 54): 'imshow_size'}, {}), '(img_lab_out, imshow_size)', True, 'import cv2 as cv\n'), ((130, 26, 130, 61), 'cv2.resize', 'cv.resize', ({(130, 36, 130, 47): 'img_bgr_out', (130, 49, 130, 60): 'imshow_size'}, {}), '(img_bgr_out, imshow_size)', True, 'import cv2 as cv\n'), ((27, 20, 27, 34), 'pathlib.Path', 'Path', ({(27, 25, 27, 33): '__file__'}, {}), '(__file__)', False, 'from pathlib import Path\n'), ((28, 20, 28, 34), 'pathlib.Path', 'Path', ({(28, 25, 28, 33): '__file__'}, {}), '(__file__)', False, 'from pathlib import Path\n')]