<reponame>jonathanyeh0723/openvino-people-counter-app
"""Import libraries for building."""
import argparse
import cv2

import os
import sys
import time

import logging as log
from inference import Network


def get_args():
    """Get command line arguments for execution."""
    parser = argparse.ArgumentParser()

    parser.add_argument("-m", "--model", required=True, type=str,
                        help="The location of model .xml file")
    parser.add_argument("-i", "--input", required=True, type=str,
                        help="The path to input stream")
    parser.add_argument("-d", "--device", default="CPU", type=str,
                        help="To specify target device for inference")
    parser.add_argument("-pt", "--prob_threshold", default=0.6, type=float,
                        help="Probability threshold for detections filtering")
    parser.add_argument("-c", "--color", default="GREEN", type=str,
                        help="The color to draw the bounding box")

    args = parser.parse_args()

    return args


def preprocessing(frame, b, c, h, w):
    """Preprocess image for inference."""
    # Resize to model's input w x h
    p_frame = cv2.resize(frame, (w, h))

    # Transpose the layout from hwc to chw
    p_frame = p_frame.transpose((2, 0, 1))

    # Reshape to model's shape b x c x h x w
    p_frame = p_frame.reshape((b, c, h, w))

    return p_frame


def infer_on_target(args):
    """Perform inference."""
    logger = log.getLogger()

    # Set default input mode
    single_image_mode = False

    # Instantiate the current request
    cur_request_id = 0

    # Initialize the plugin
    plugin = Network()

    # Load the model and get the input data
    plugin.load_model(args.model, args.device, cur_request_id)
    b, c, h, w = plugin.get_input_shape()

    begin_time = time.time()

    if args.input == 'CAM':
        input_stream = 0

    elif args.input.endswith('.jpg') or args.input.endswith('.bmp'):
        single_image_mode = True
        input_stream = args.input

    else:
        input_stream = args.input
        assert os.path.isfile(args.input), "Specified input file doesn't exist"

    # Create VideoCapture object, and have it opened
    cap = cv2.VideoCapture(input_stream)

    if not cap.isOpened():
        log.error("ERROR! Unable to open the video source")

    # Grab the dimension of input width and height
    global initial_w, initial_h
    initial_w = cap.get(3)
    initial_h = cap.get(4)

    # Set initial parameters for counting
    last_count = 0
    total_count = 0
    frame_count = 0
    detected_frame = 0
    not_counted = 0

    # Iterate the frames until the video ends
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        key_pressed = cv2.waitKey(60)
        frame_count += 1

        if not single_image_mode:
            total_info = "Total people counted: {}".format(total_count - 1)
            cv2.putText(frame, total_info, (525, 15), cv2.FONT_HERSHEY_COMPLEX,
                        0.5, (200, 10, 10), 1)

        # Preprocessing input data
        p_frame = preprocessing(frame, b, c, h, w)

        # Perform async inference
        plugin.exec_net(cur_request_id, p_frame)
        inf_start = time.time()

        # Get the inference result
        if plugin.wait(cur_request_id) == 0:
            result = plugin.extract_output(cur_request_id)
            det_time = time.time() - inf_start

            # Draw the bounding boxes
            frame, current_count, flag = draw_boxes(frame, result,
                                                    args, not_counted)

            counted_info = "Person detected: "
            cv2.putText(frame, counted_info, (525, 35),
                        cv2.FONT_HERSHEY_COMPLEX, 0.5, (200, 10, 10), 1)

            if flag is True:
                detected_frame += 1
                not_counted = 0
                cv2.putText(frame, "{}".format(flag), (680, 35),
                            cv2.FONT_HERSHEY_COMPLEX, 0.5, (0, 255, 0), 1)

            if flag is False:
                detected_frame = 0
                not_counted += 1
                cv2.putText(frame, "{}".format(flag), (680, 35),
                            cv2.FONT_HERSHEY_COMPLEX, 0.5, (0, 0, 255), 1)

            inf_info = "Inference time: {:.3f}ms".format(det_time * 1000)
            cv2.putText(frame, inf_info, (15, 15), cv2.FONT_HERSHEY_COMPLEX,
                        0.5, (200, 10, 10), 1)

            frame_info = "Total frames counted: {}".format(frame_count)
            cv2.putText(frame, frame_info, (15, 35), cv2.FONT_HERSHEY_COMPLEX,
                        0.5, (200, 10, 10), 1)

            check_info = "Number of frames when person is detected: {}".format(detected_frame)
            cv2.putText(frame, check_info, (15, 75), cv2.FONT_HERSHEY_COMPLEX,
                        0.5, (200, 10, 10), 1)

            lost_info = "Lost frame from last counting: {}".format(not_counted)
            cv2.putText(frame, lost_info, (15, 95), cv2.FONT_HERSHEY_COMPLEX,
                        0.5, (200, 10, 10), 1)

            if current_count > last_count:
                start = time.time()
                total_count = total_count + current_count - last_count

            if current_count < last_count:
                duration = int(time.time() - start)

            last_count = current_count

            if key_pressed == 27:
                break

        cv2.imshow('demo', frame)

        # If single image mode, write out the inference result
        if single_image_mode:
            cv2.imwrite('./output_image.jpeg', frame)

    cv2.waitKey(0)

    # Release all. Destroy all OpenCV windows.
    cap.release()
    cv2.destroyAllWindows()

    total_time = time.time() - begin_time
    logger.error("[ INFO ] Conducted program successfully.")
    logger.error("[ INFO ] Total frames counted: {}".format(frame_count))
    logger.error("[ INFO ] Total execution time: {:.3f} seconds".format(total_time))


def draw_boxes(frame, result, args, not_counted):
    """Draw the bounding boxes if person is detected."""
    current_count = 0
    detect_flag = False

    color_dict = {"BLUE": (255, 0, 0), "GREEN": (0, 255, 0),
                  "RED": (0, 0, 255)}
    if args.color:
        out_color = color_dict.get(args.color)
    else:
        out_color = color_dict["GREEN"]

    for obj in result[0][0]:
        conf = obj[2]
        if conf >= args.prob_threshold:
            xmin = int(obj[3] * initial_w)
            ymin = int(obj[4] * initial_h)
            xmax = int(obj[5] * initial_w)
            ymax = int(obj[6] * initial_h)

            current_count += 1
            cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), out_color, 3)
            detect_flag = True

        if detect_flag is False and not_counted < 15:
            current_count = 1

    return frame, current_count, detect_flag


def main():
    """Put all together. To execute the app."""
    args = get_args()
    infer_on_target(args)


if __name__ == "__main__":
    main()

"""Import libraries for building."""
import argparse
import cv2

import os
import sys
import time

import logging as log
from inference import Network


def get_args():
    """Get command line arguments for execution."""
    parser = argparse.ArgumentParser()

    parser.add_argument("-m", "--model", required=True, type=str,
                        help="The location of model .xml file")
    parser.add_argument("-i", "--input", required=True, type=str,
                        help="The path to input stream")
    parser.add_argument("-d", "--device", default="CPU", type=str,
                        help="To specify target device for inference")
    parser.add_argument("-pt", "--prob_threshold", default=0.6, type=float,
                        help="Probability threshold for detections filtering")
    parser.add_argument("-c", "--color", default="GREEN", type=str,
                        help="The color to draw the bounding box")

    args = parser.parse_args()

    return args


def preprocessing(frame, b, c, h, w):
    """Preprocess image for inference."""
    # Resize to model's input w x h
    p_frame = cv2.resize(frame, (w, h))

    # Transpose the layout from hwc to chw
    p_frame = p_frame.transpose((2, 0, 1))

    # Reshape to model's shape b x c x h x w
    p_frame = p_frame.reshape((b, c, h, w))

    return p_frame


def infer_on_target(args):
    """Perform inference."""
    logger = log.getLogger()

    # Set default input mode
    single_image_mode = False

    # Instantiate the current request
    cur_request_id = 0

    # Initialize the plugin
    plugin = Network()

    # Load the model and get the input data
    plugin.load_model(args.model, args.device, cur_request_id)
    b, c, h, w = plugin.get_input_shape()

    begin_time = time.time()

    if args.input == 'CAM':
        input_stream = 0

    elif args.input.endswith('.jpg') or args.input.endswith('.bmp'):
        single_image_mode = True
        input_stream = args.input

    else:
        input_stream = args.input
        assert os.path.isfile(args.input), "Specified input file doesn't exist"

    # Create VideoCapture object, and have it opened
    cap = cv2.VideoCapture(input_stream)

    if not cap.isOpened():
        log.error("ERROR! Unable to open the video source")

    # Grab the dimension of input width and height
    global initial_w, initial_h
    initial_w = cap.get(3)
    initial_h = cap.get(4)

    # Set initial parameters for counting
    last_count = 0
    total_count = 0
    frame_count = 0
    detected_frame = 0
    not_counted = 0

    # Iterate the frames until the video ends
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        key_pressed = cv2.waitKey(60)
        frame_count += 1

        if not single_image_mode:
            total_info = "Total people counted: {}".format(total_count - 1)
            cv2.putText(frame, total_info, (525, 15), cv2.FONT_HERSHEY_COMPLEX,
                        0.5, (200, 10, 10), 1)

        # Preprocessing input data
        p_frame = preprocessing(frame, b, c, h, w)

        # Perform async inference
        plugin.exec_net(cur_request_id, p_frame)
        inf_start = time.time()

        # Get the inference result
        if plugin.wait(cur_request_id) == 0:
            result = plugin.extract_output(cur_request_id)
            det_time = time.time() - inf_start

            # Draw the bounding boxes
            frame, current_count, flag = draw_boxes(frame, result,
                                                    args, not_counted)

            counted_info = "Person detected: "
            cv2.putText(frame, counted_info, (525, 35),
                        cv2.FONT_HERSHEY_COMPLEX, 0.5, (200, 10, 10), 1)

            if flag is True:
                detected_frame += 1
                not_counted = 0
                cv2.putText(frame, "{}".format(flag), (680, 35),
                            cv2.FONT_HERSHEY_COMPLEX, 0.5, (0, 255, 0), 1)

            if flag is False:
                detected_frame = 0
                not_counted += 1
                cv2.putText(frame, "{}".format(flag), (680, 35),
                            cv2.FONT_HERSHEY_COMPLEX, 0.5, (0, 0, 255), 1)

            inf_info = "Inference time: {:.3f}ms".format(det_time * 1000)
            cv2.putText(frame, inf_info, (15, 15), cv2.FONT_HERSHEY_COMPLEX,
                        0.5, (200, 10, 10), 1)

            frame_info = "Total frames counted: {}".format(frame_count)
            cv2.putText(frame, frame_info, (15, 35), cv2.FONT_HERSHEY_COMPLEX,
                        0.5, (200, 10, 10), 1)

            check_info = "Number of frames when person is detected: {}".format(detected_frame)
            cv2.putText(frame, check_info, (15, 75), cv2.FONT_HERSHEY_COMPLEX,
                        0.5, (200, 10, 10), 1)

            lost_info = "Lost frame from last counting: {}".format(not_counted)
            cv2.putText(frame, lost_info, (15, 95), cv2.FONT_HERSHEY_COMPLEX,
                        0.5, (200, 10, 10), 1)

            if current_count > last_count:
                start = time.time()
                total_count = total_count + current_count - last_count

            if current_count < last_count:
                duration = int(time.time() - start)

            last_count = current_count

            if key_pressed == 27:
                break

        cv2.imshow('demo', frame)

        # If single image mode, write out the inference result
        if single_image_mode:
            cv2.imwrite('./output_image.jpeg', frame)

    cv2.waitKey(0)

    # Release all. Destroy all OpenCV windows.
    cap.release()
    cv2.destroyAllWindows()

    total_time = time.time() - begin_time
    logger.error("[ INFO ] Conducted program successfully.")
    logger.error("[ INFO ] Total frames counted: {}".format(frame_count))
    logger.error("[ INFO ] Total execution time: {:.3f} seconds".format(total_time))


def draw_boxes(frame, result, args, not_counted):
    """Draw the bounding boxes if person is detected."""
    current_count = 0
    detect_flag = False

    color_dict = {"BLUE": (255, 0, 0), "GREEN": (0, 255, 0),
                  "RED": (0, 0, 255)}
    if args.color:
        out_color = color_dict.get(args.color)
    else:
        out_color = color_dict["GREEN"]

    for obj in result[0][0]:
        conf = obj[2]
        if conf >= args.prob_threshold:
            xmin = int(obj[3] * initial_w)
            ymin = int(obj[4] * initial_h)
            xmax = int(obj[5] * initial_w)
            ymax = int(obj[6] * initial_h)

            current_count += 1
            cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), out_color, 3)
            detect_flag = True

        if detect_flag is False and not_counted < 15:
            current_count = 1

    return frame, current_count, detect_flag


def main():
    """Put all together. To execute the app."""
    args = get_args()
    infer_on_target(args)


if __name__ == "__main__":
    main()
Import libraries for building. Get command line arguments for execution. Preprocess image for inference. # Resize to model's input w x h # Transpose the layout from hwc to chw # Reshape to model's shape b x c x h x w Perform inference. # Set default input mode # Instantiate the current request # Initialize the plugin # Load the model and get the input data # Create VideoCapture object, and have it opened # Grab the dimension of input width and height # Set initial parameters for counting # Iterate the frames until the video ends # Preprocessing input data # Perform async inference # Get the inference result # Draw the bounding boxes # If single image mode, write out the inference result # Release all. Destroy all OpenCV windows. Draw the bounding boxes if person is detected. Put all together. To execute the app.