what.examples.yolov4_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.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():
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)