what.examples.yolov3_demo

 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.yolov3 import YOLOV3
 8from what.models.detection.yolo.yolov3_tiny import YOLOV3_TINY
 9
10from what.cli.model import *
11
12from what.utils.file import get_file
13
14what_yolov3_model_list = what_model_list[0:4]
15
16def yolov3_inference_demo():
17
18    max_len = max([len(x[WHAT_MODEL_NAME_INDEX]) for x in what_yolov3_model_list])
19    for i, model in enumerate(what_yolov3_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_yolov3_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_yolov3_model_list[index][WHAT_MODEL_FILE_INDEX])):
34        get_file(what_yolov3_model_list[index][WHAT_MODEL_FILE_INDEX],
35                 WHAT_MODEL_PATH,
36                 what_yolov3_model_list[index][WHAT_MODEL_URL_INDEX],
37                 what_yolov3_model_list[index][WHAT_MODEL_HASH_INDEX])
38
39    if index == 0 or index == 1:
40        model = YOLOV3(COCO_CLASS_NAMES, os.path.join(WHAT_MODEL_PATH, what_yolov3_model_list[index][WHAT_MODEL_FILE_INDEX]))
41
42    if index == 2 or index == 3:
43        model = YOLOV3_TINY(COCO_CLASS_NAMES, os.path.join(WHAT_MODEL_PATH, what_yolov3_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('YOLOv3 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 yolov3_inference_demo():
17def yolov3_inference_demo():
18
19    max_len = max([len(x[WHAT_MODEL_NAME_INDEX]) for x in what_yolov3_model_list])
20    for i, model in enumerate(what_yolov3_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_yolov3_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_yolov3_model_list[index][WHAT_MODEL_FILE_INDEX])):
35        get_file(what_yolov3_model_list[index][WHAT_MODEL_FILE_INDEX],
36                 WHAT_MODEL_PATH,
37                 what_yolov3_model_list[index][WHAT_MODEL_URL_INDEX],
38                 what_yolov3_model_list[index][WHAT_MODEL_HASH_INDEX])
39
40    if index == 0 or index == 1:
41        model = YOLOV3(COCO_CLASS_NAMES, os.path.join(WHAT_MODEL_PATH, what_yolov3_model_list[index][WHAT_MODEL_FILE_INDEX]))
42
43    if index == 2 or index == 3:
44        model = YOLOV3_TINY(COCO_CLASS_NAMES, os.path.join(WHAT_MODEL_PATH, what_yolov3_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('YOLOv3 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)