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    WHAT_YOLOV4_MODEL_FILE = what_yolov4_model_list[index][WHAT_MODEL_FILE_INDEX]
34    WHAT_YOLOV4_MODEL_URL  = what_yolov4_model_list[index][WHAT_MODEL_URL_INDEX]
35    WHAT_YOLOV4_MODEL_HASH = what_yolov4_model_list[index][WHAT_MODEL_HASH_INDEX]
36
37    if not os.path.isfile(os.path.join(WHAT_MODEL_PATH, WHAT_YOLOV4_MODEL_FILE)):
38        get_file(WHAT_YOLOV4_MODEL_FILE,
39                WHAT_MODEL_PATH,
40                WHAT_YOLOV4_MODEL_URL,
41                WHAT_YOLOV4_MODEL_HASH)
42
43    if index == 0:
44        model = YOLOV4(COCO_CLASS_NAMES, os.path.join(WHAT_MODEL_PATH, WHAT_YOLOV4_MODEL_FILE))
45
46    if index == 1:
47        model = YOLOV4_TINY(COCO_CLASS_NAMES, os.path.join(WHAT_MODEL_PATH, WHAT_YOLOV4_MODEL_FILE))
48
49    video = input(f"Please input the OpenCV capture device (e.g. 0, 1, 2): ")
50
51    while not video.isdigit():
52        video = input(f"Please input the OpenCV capture device (e.g. 0, 1, 2): ")
53
54    try:
55        # Capture from camera or video
56        if video.isdigit():
57            cap = cv2.VideoCapture(int(video))
58        else:
59            cap = cv2.VideoCapture(video)
60
61        #cap.set(3, 1920)
62        #cap.set(4, 1080)
63
64        while True:
65            _, orig_image = cap.read()
66            if orig_image is None:
67                continue
68
69            # Image preprocessing
70            image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
71
72            # Run inference
73            images, boxes, labels, probs = model.predict(image)
74            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
75
76            # Draw bounding boxes onto the image
77            if len(boxes) > 0:
78                output = draw_bounding_boxes(image, boxes, labels, model.class_names, probs);
79
80            cv2.imshow('YOLOv4 Demo', image)
81
82            if cv2.waitKey(1) & 0xFF == ord('q'):
83                break
84
85        cap.release()
86        cv2.destroyAllWindows()
87
88    except Exception as e:
89        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    WHAT_YOLOV4_MODEL_FILE = what_yolov4_model_list[index][WHAT_MODEL_FILE_INDEX]
35    WHAT_YOLOV4_MODEL_URL  = what_yolov4_model_list[index][WHAT_MODEL_URL_INDEX]
36    WHAT_YOLOV4_MODEL_HASH = what_yolov4_model_list[index][WHAT_MODEL_HASH_INDEX]
37
38    if not os.path.isfile(os.path.join(WHAT_MODEL_PATH, WHAT_YOLOV4_MODEL_FILE)):
39        get_file(WHAT_YOLOV4_MODEL_FILE,
40                WHAT_MODEL_PATH,
41                WHAT_YOLOV4_MODEL_URL,
42                WHAT_YOLOV4_MODEL_HASH)
43
44    if index == 0:
45        model = YOLOV4(COCO_CLASS_NAMES, os.path.join(WHAT_MODEL_PATH, WHAT_YOLOV4_MODEL_FILE))
46
47    if index == 1:
48        model = YOLOV4_TINY(COCO_CLASS_NAMES, os.path.join(WHAT_MODEL_PATH, WHAT_YOLOV4_MODEL_FILE))
49
50    video = input(f"Please input the OpenCV capture device (e.g. 0, 1, 2): ")
51
52    while not video.isdigit():
53        video = input(f"Please input the OpenCV capture device (e.g. 0, 1, 2): ")
54
55    try:
56        # Capture from camera or video
57        if video.isdigit():
58            cap = cv2.VideoCapture(int(video))
59        else:
60            cap = cv2.VideoCapture(video)
61
62        #cap.set(3, 1920)
63        #cap.set(4, 1080)
64
65        while True:
66            _, orig_image = cap.read()
67            if orig_image is None:
68                continue
69
70            # Image preprocessing
71            image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
72
73            # Run inference
74            images, boxes, labels, probs = model.predict(image)
75            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
76
77            # Draw bounding boxes onto the image
78            if len(boxes) > 0:
79                output = draw_bounding_boxes(image, boxes, labels, model.class_names, probs);
80
81            cv2.imshow('YOLOv4 Demo', image)
82
83            if cv2.waitKey(1) & 0xFF == ord('q'):
84                break
85
86        cap.release()
87        cv2.destroyAllWindows()
88
89    except Exception as e:
90        print(e)