what.examples.yolov4_demo API documentation

# what.examples.yolov4_demo

View Source

     1import cv2
     2import os.path
     3
     4from what.models.detection.datasets.coco import COCO_CLASS_NAMES
     5from what.models.detection.utils.box_utils import draw_bounding_boxes
     6
     7from what.models.detection.yolo.yolov4 import YOLOV4
     8from what.models.detection.yolo.yolov4_tiny import YOLOV4_TINY
     9
    10from what.cli.model import *
    11
    12from what.utils.file import get_file
    13
    14what_yolov4_model_list = what_model_list[4:6]
    15
    16def yolov4_inference_demo():
    17
    18    max_len = max([len(x[WHAT_MODEL_NAME_INDEX]) for x in what_yolov4_model_list])
    19    for i, model in enumerate(what_yolov4_model_list, start=1):
    20        if os.path.isfile(os.path.join(WHAT_MODEL_PATH, model[WHAT_MODEL_FILE_INDEX])):
    21            downloaded = 'x'
    22        else:
    23            downloaded = ' '
    24        print('[{}] {} : {:<{w}s}\t{}\t{}'.format(downloaded, i, model[WHAT_MODEL_NAME_INDEX], model[WHAT_MODEL_TYPE_INDEX], model[WHAT_MODEL_DESC_INDEX], w=max_len))
    25
    26    index = input(f"Please input the model index: ")
    27    while not index.isdigit() or int(index) > len(what_yolov4_model_list):
    28        index = input(f"Model [{index}] does not exist. Please try again: ")
    29
    30    index = int(index) - 1
    31
    32    # Download the model first if not exists
    33    if not os.path.isfile(os.path.join(WHAT_MODEL_PATH, what_yolov4_model_list[index][WHAT_MODEL_FILE_INDEX])):
    34        get_file(what_yolov4_model_list[index][WHAT_MODEL_FILE_INDEX],
    35                 WHAT_MODEL_PATH,
    36                 what_yolov4_model_list[index][WHAT_MODEL_URL_INDEX],
    37                 what_yolov4_model_list[index][WHAT_MODEL_HASH_INDEX])
    38
    39    if index == 0:
    40        model = YOLOV4(COCO_CLASS_NAMES, os.path.join(WHAT_MODEL_PATH, what_yolov4_model_list[index][WHAT_MODEL_FILE_INDEX]))
    41
    42    if index == 1:
    43        model = YOLOV4_TINY(COCO_CLASS_NAMES, os.path.join(WHAT_MODEL_PATH, what_yolov4_model_list[index][WHAT_MODEL_FILE_INDEX]))
    44
    45    video = input(f"Please input the OpenCV capture device (e.g. 0, 1, 2): ")
    46
    47    while not video.isdigit():
    48        video = input(f"Please input the OpenCV capture device (e.g. 0, 1, 2): ")
    49
    50    try:
    51        # Capture from camera or video
    52        if video.isdigit():
    53            cap = cv2.VideoCapture(int(video))
    54        else:
    55            cap = cv2.VideoCapture(video)
    56
    57        #cap.set(3, 1920)
    58        #cap.set(4, 1080)
    59
    60        while True:
    61            _, orig_image = cap.read()
    62            if orig_image is None:
    63                continue
    64
    65            # Image preprocessing
    66            image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
    67
    68            # Run inference
    69            images, boxes, labels, probs = model.predict(image, 10, 0.4)
    70            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    71
    72            # Draw bounding boxes onto the image
    73            output = draw_bounding_boxes(image, boxes, labels, model.class_names, probs);
    74
    75            cv2.imshow('YOLOv4 Demo', image)
    76
    77            if cv2.waitKey(1) & 0xFF == ord('q'):
    78                break
    79
    80        cap.release()
    81        cv2.destroyAllWindows()
    82
    83    except Exception as e:
    84        print(e)

def yolov4_inference_demo(): View Source

    17def yolov4_inference_demo():
    18
    19    max_len = max([len(x[WHAT_MODEL_NAME_INDEX]) for x in what_yolov4_model_list])
    20    for i, model in enumerate(what_yolov4_model_list, start=1):
    21        if os.path.isfile(os.path.join(WHAT_MODEL_PATH, model[WHAT_MODEL_FILE_INDEX])):
    22            downloaded = 'x'
    23        else:
    24            downloaded = ' '
    25        print('[{}] {} : {:<{w}s}\t{}\t{}'.format(downloaded, i, model[WHAT_MODEL_NAME_INDEX], model[WHAT_MODEL_TYPE_INDEX], model[WHAT_MODEL_DESC_INDEX], w=max_len))
    26
    27    index = input(f"Please input the model index: ")
    28    while not index.isdigit() or int(index) > len(what_yolov4_model_list):
    29        index = input(f"Model [{index}] does not exist. Please try again: ")
    30
    31    index = int(index) - 1
    32
    33    # Download the model first if not exists
    34    if not os.path.isfile(os.path.join(WHAT_MODEL_PATH, what_yolov4_model_list[index][WHAT_MODEL_FILE_INDEX])):
    35        get_file(what_yolov4_model_list[index][WHAT_MODEL_FILE_INDEX],
    36                 WHAT_MODEL_PATH,
    37                 what_yolov4_model_list[index][WHAT_MODEL_URL_INDEX],
    38                 what_yolov4_model_list[index][WHAT_MODEL_HASH_INDEX])
    39
    40    if index == 0:
    41        model = YOLOV4(COCO_CLASS_NAMES, os.path.join(WHAT_MODEL_PATH, what_yolov4_model_list[index][WHAT_MODEL_FILE_INDEX]))
    42
    43    if index == 1:
    44        model = YOLOV4_TINY(COCO_CLASS_NAMES, os.path.join(WHAT_MODEL_PATH, what_yolov4_model_list[index][WHAT_MODEL_FILE_INDEX]))
    45
    46    video = input(f"Please input the OpenCV capture device (e.g. 0, 1, 2): ")
    47
    48    while not video.isdigit():
    49        video = input(f"Please input the OpenCV capture device (e.g. 0, 1, 2): ")
    50
    51    try:
    52        # Capture from camera or video
    53        if video.isdigit():
    54            cap = cv2.VideoCapture(int(video))
    55        else:
    56            cap = cv2.VideoCapture(video)
    57
    58        #cap.set(3, 1920)
    59        #cap.set(4, 1080)
    60
    61        while True:
    62            _, orig_image = cap.read()
    63            if orig_image is None:
    64                continue
    65
    66            # Image preprocessing
    67            image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
    68
    69            # Run inference
    70            images, boxes, labels, probs = model.predict(image, 10, 0.4)
    71            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    72
    73            # Draw bounding boxes onto the image
    74            output = draw_bounding_boxes(image, boxes, labels, model.class_names, probs);
    75
    76            cv2.imshow('YOLOv4 Demo', image)
    77
    78            if cv2.waitKey(1) & 0xFF == ord('q'):
    79                break
    80
    81        cap.release()
    82        cv2.destroyAllWindows()
    83
    84    except Exception as e:
    85        print(e)