what.examples.yolox_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.yolox.yolox_x import YOLOX_X
  8from what.models.detection.yolox.yolox_l import YOLOX_L
  9from what.models.detection.yolox.yolox_m import YOLOX_M
 10from what.models.detection.yolox.yolox_s import YOLOX_S
 11
 12from what.cli.model import *
 13
 14from what.utils.file import get_file
 15
 16what_yolox_model_list = what_model_list[9:13]
 17
 18def yolox_inference_demo():
 19
 20    max_len = max([len(x[WHAT_MODEL_NAME_INDEX]) for x in what_yolox_model_list])
 21    for i, model in enumerate(what_yolox_model_list, start=1):
 22        if os.path.isfile(os.path.join(WHAT_MODEL_PATH, model[WHAT_MODEL_FILE_INDEX])):
 23            downloaded = 'x'
 24        else:
 25            downloaded = ' '
 26        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))
 27
 28    index = input(f"Please input the model index: ")
 29    while not index.isdigit() or int(index) > len(what_yolox_model_list):
 30        index = input(f"Model [{index}] does not exist. Please try again: ")
 31
 32    index = int(index) - 1
 33
 34    # Download the model first if not exists
 35    WHAT_YOLOX_MODEL_FILE = what_yolox_model_list[index][WHAT_MODEL_FILE_INDEX]
 36    WHAT_YOLOX_MODEL_URL  = what_yolox_model_list[index][WHAT_MODEL_URL_INDEX]
 37    WHAT_YOLOX_MODEL_HASH = what_yolox_model_list[index][WHAT_MODEL_HASH_INDEX]
 38
 39    if not os.path.isfile(os.path.join(WHAT_MODEL_PATH, WHAT_YOLOX_MODEL_FILE)):
 40        get_file(WHAT_YOLOX_MODEL_FILE,
 41                WHAT_MODEL_PATH,
 42                WHAT_YOLOX_MODEL_URL,
 43                WHAT_YOLOX_MODEL_HASH)
 44
 45    if index == 0:
 46        model = YOLOX_X(COCO_CLASS_NAMES, os.path.join(WHAT_MODEL_PATH, WHAT_YOLOX_MODEL_FILE))
 47
 48    if index == 1:
 49        model = YOLOX_L(COCO_CLASS_NAMES, os.path.join(WHAT_MODEL_PATH, WHAT_YOLOX_MODEL_FILE))
 50
 51    if index == 2:
 52        model = YOLOX_M(COCO_CLASS_NAMES, os.path.join(WHAT_MODEL_PATH, WHAT_YOLOX_MODEL_FILE))
 53
 54    if index == 3:
 55        model = YOLOX_S(COCO_CLASS_NAMES, os.path.join(WHAT_MODEL_PATH, WHAT_YOLOX_MODEL_FILE))
 56
 57    video = input(f"Please input the OpenCV capture device (e.g. 0, 1, 2): ")
 58
 59    while not video.isdigit():
 60        video = input(f"Please input the OpenCV capture device (e.g. 0, 1, 2): ")
 61
 62    try:
 63        # Capture from camera or video
 64        if video.isdigit():
 65            cap = cv2.VideoCapture(int(video))
 66        else:
 67            cap = cv2.VideoCapture(video)
 68
 69        #cap.set(3, 1920)
 70        #cap.set(4, 1080)
 71
 72        while True:
 73            _, orig_image = cap.read()
 74            if orig_image is None:
 75                continue
 76
 77            # Image preprocessing
 78            image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
 79
 80            # Run inference
 81            images, boxes, labels, probs = model.predict(image)
 82            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
 83
 84            # Draw bounding boxes onto the image
 85            if len(boxes) > 0:
 86                output = draw_bounding_boxes(image, boxes, labels, model.class_names, probs);
 87
 88            cv2.imshow('YOLOX Demo', image)
 89
 90            if cv2.waitKey(1) & 0xFF == ord('q'):
 91                break
 92
 93        cap.release()
 94        cv2.destroyAllWindows()
 95
 96    except Exception as e:
 97        print(e)
 98
 99
100if __name__ == "__main__":
101    yolox_inference_demo()