main.py

from core.posemodule import poseDetector
from core.analysis import create_dashboard
from core.analysis import analysis
from centroid_tracking.tracker import Tracker
from core.utils import *
from tensorflow.python.keras.models import load_model
from tensorflow.compat.v1 import InteractiveSession
from tensorflow.compat.v1 import ConfigProto
from absl.flags import FLAGS
import numpy as np
import cv2
from absl import app, flags
from tensorflow.python.saved_model import tag_constants
import tensorflow as tf
import time
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"


flags.DEFINE_string('framework', 'tf', '(tf, tflite, trt')
flags.DEFINE_string('weights_ball', './checkpoints/3l4b3_ball_416',
                    'path to weights ball file')
flags.DEFINE_string('weights_palm', './checkpoints/custom-tiny-palm-416',
                    'path to weights palm file')
flags.DEFINE_integer('size', 416, 'resize images to')
flags.DEFINE_boolean('tiny', True, 'yolo or yolo-tiny')
flags.DEFINE_string('model', 'yolov4-tiny-3l', 'yolov3 or yolov4')
flags.DEFINE_string('video', 'src/video0.mov', 'path to input video')
flags.DEFINE_float('iou', 0.25, 'iou threshold')
flags.DEFINE_float('score', 0.30, 'score threshold')
flags.DEFINE_string('output_format', 'XVID',
                    'codec used in VideoWriter when saving video to file')
flags.DEFINE_string('output', 'output.avi', 'path to output video')
flags.DEFINE_string('demo_output', 'demo.avi', 'path to demo output video')
flags.DEFINE_string('ptn_model', 'checkpoints/pattern_model.h5',
                    'path to pattern recognition model')
flags.DEFINE_boolean('gpu', True, 'activate gpu - True else False')


def main(_argv):
    # initialize all the FLAGS setting
    input_size = FLAGS.size
    video_path = FLAGS.video
    gpu = FLAGS.gpu
    weights_ball = FLAGS.weights_ball
    weights_palm = FLAGS.weights_palm
    score = FLAGS.score
    iou = FLAGS.iou
    pattern_model = FLAGS.ptn_model
    output = FLAGS.output
    demo_output = FLAGS.demo_output
    output_format = FLAGS.output_format

    if gpu:
        # set up gpu setting
        physical_devices = tf.config.experimental.list_physical_devices('GPU')
        if len(physical_devices) > 0:
            tf.config.experimental.set_memory_growth(physical_devices[0], True)
        config = ConfigProto()
        config.gpu_options.allow_growth = True
        session = InteractiveSession(config=config)

    # load in all the models
    saved_model_loaded_ball = tf.saved_model.load(
        weights_ball, tags=[tag_constants.SERVING])
    infer_ball = saved_model_loaded_ball.signatures['serving_default']
    pattern_model = load_model(pattern_model)
    pose_detector = poseDetector()  # human pose estimator

    # read in the video
    print("Video from: ", video_path)
    try:
        vid = cv2.VideoCapture(int(video_path))  # 0 - real time camera access
    except:
        vid = cv2.VideoCapture(video_path)  # else - video input

    # get os resolution for display purpose
    rescale_width, rescale_height, image_width, image_height = resolution_display(
        vid)

    # initialize video writer
    if output:
        fps = 20
        codec = cv2.VideoWriter_fourcc(*output_format)
        out = cv2.VideoWriter(output, codec, fps, (image_width, image_height))
        demo_out = cv2.VideoWriter(
            demo_output, codec, fps, (image_width, image_height))

    tracker = Tracker()  # initialize tracker
    ptns = []

    # start capturing and detection
    while True:
        return_value, frame = vid.read()
        if not return_value:
            print("Video processing complete")
            os._exit(0)

        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frame = cv2.resize(frame, (image_width, image_height))
        demo = np.zeros((image_height, image_width, 3), np.uint8)
        prev_time = time.time()  # fps counting

        # human pose estimation
        sucess, demo = pose_detector.findPose(frame, demo)
        # run the detection only if human pose found
        if sucess:
            lmList = pose_detector.findPosition(demo, draw=False)
            right_palm, left_palm = pose_detector.findPalm()

            # ball detection
            image_data = cv2.resize(frame, (input_size, input_size))
            image_data = image_data / 255.
            image_data = image_data[np.newaxis, ...].astype(np.float32)

            # capture the detection box
            batch_data = tf.constant(image_data)
            pred_bbox_ball = infer_ball(batch_data)
            for key, value in pred_bbox_ball.items():
                boxes_ball = value[:, :, 0:4]
                pred_conf_ball = value[:, :, 4:]

            # non max suppression
            boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression(
                boxes=tf.reshape(
                    boxes_ball, (tf.shape(boxes_ball)[0], -1, 1, 4)),
                scores=tf.reshape(
                    pred_conf_ball, (tf.shape(pred_conf_ball)[0], -1, tf.shape(pred_conf_ball)[-1])),
                max_output_size_per_class=50,
                max_total_size=50,
                iou_threshold=iou,
                score_threshold=score
            )

            # finalized pred bbox
            pred_bbox = [boxes.numpy(), scores.numpy(
            ), classes.numpy(), valid_detections.numpy()]

            # perform unbound tracking
            pair_ball = tracker.track(frame, pred_bbox)
            bound_ball = mapping(pair_ball, [right_palm, left_palm], True)
            tracker.object_checking(bound_ball)
            bound_ball = mapping(tracker.pair_ball, [right_palm, left_palm])
            bound_ball_copy = copy.deepcopy(bound_ball)
            unbound_results = classification(
                frame, bound_ball, tracker.prev_pair_ball, tracker.pair_ball, pattern_model)

            # analysis result and display on dashboard
            frame = create_dashboard(frame)
            frame = analysis(pair_ball, tracker.pair_ball, frame)
            frame = pose_detector.distance_estimation(frame)
            frame, dmeo = pose_detector.find_Elbow_angle(frame, demo)

            # perform bound tracking
            demo, pred_balls = tracker.bound_tracking(
                demo, unbound_results, bound_ball_copy)
            bound_results = classification(
                frame, pred_balls, tracker.prev_pair_ball, tracker.pair_ball, pattern_model)
            results = unbound_results + bound_results
            bound_ball.extend(pred_balls)

            # display result - simulation and draw bbox on frame
            demo, ptns = display_demo(demo, results, ptns, bound_ball, tracker.pair_ball, [
                                      right_palm, left_palm])
            frame = draw_bbox(frame, bound_ball, tracker.pair_ball, [
                              right_palm, left_palm])

            # display frame
            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            re_frame = cv2.resize(frame, (rescale_width, rescale_height))
            re_demo = cv2.resize(demo, (rescale_width, rescale_height))

        # set the position for output and demo result window
        if sucess:
            cv2.imshow("output", re_frame)
            cv2.imshow("demo", re_demo)
        else:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            cv2.imshow("output", frame)
            cv2.imshow("demo", demo)
        cv2.moveWindow("output", 0, 0)
        cv2.moveWindow("demo", int(rescale_width), 0)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

        # printout fps
        curr_time = time.time()
        exec_time = 1.0 / (curr_time - prev_time)
        print("FPS: %.2f" % exec_time)
        print()
        print()

        # save both output and demo results
        if output:
            out.write(frame)

        if demo_output:
            demo_out.write(demo)


if __name__ == '__main__':
    try:
        app.run(main)
    except SystemExit:
        pass