what.examples.faster_rcnn_demo API documentation

# what.examples.faster_rcnn_demo

View Source

     1import cv2
     2import torch
     3import numpy as np
     4
     5from what.models.detection.frcnn.faster_rcnn import FasterRCNN
     6# from what.models.detection.frcnn.datasets.util import read_image
     7
     8from what.cli.model import *
     9from what.utils.file import get_file
    10
    11from what.models.detection.utils.box_utils import draw_bounding_boxes
    12from what.models.detection.datasets.voc import VOC_CLASS_NAMES
    13
    14def frcnn_inference_demo():
    15
    16    video = input(f"Please input the OpenCV capture device (e.g. 0, 1, 2): ")
    17
    18    while not video.isdigit():
    19        video = input(f"Please input the OpenCV capture device (e.g. 0, 1, 2): ")
    20
    21    # Capture from camera
    22    cap = cv2.VideoCapture(int(video))
    23    #cap.set(3, 1920)
    24    #cap.set(4, 1080)
    25
    26    # Download the model first if not exists
    27    # Check what_model_list for all available models
    28    index = 8
    29    if not os.path.isfile(os.path.join(WHAT_MODEL_PATH, what_model_list[index][WHAT_MODEL_FILE_INDEX])):
    30        get_file(what_model_list[index][WHAT_MODEL_FILE_INDEX],
    31                 WHAT_MODEL_PATH,
    32                 what_model_list[index][WHAT_MODEL_URL_INDEX],
    33                 what_model_list[index][WHAT_MODEL_HASH_INDEX])
    34
    35    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    36
    37    model = FasterRCNN(device=device)
    38
    39    model.load(os.path.join(WHAT_MODEL_PATH, what_model_list[index][WHAT_MODEL_FILE_INDEX]), map_location=device)
    40
    41    while True:
    42        _, orig_image = cap.read()
    43        if orig_image is None:
    44            continue
    45
    46        # Image preprocessing
    47        input_img = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
    48        height, width, _ = input_img.shape
    49
    50        # NHWC -> NCHW
    51        input_img = np.array(input_img).transpose((2, 0, 1))
    52        input_img = torch.from_numpy(input_img)[None]
    53
    54        # img = read_image('notebooks/demo.jpg', format='NCHW')
    55        # input_img = torch.from_numpy(img)[None]
    56        # RGB --> BGR
    57        # img = img.transpose((1, 2, 0)).astype(np.uint8)
    58
    59        inputs, boxes, labels, scores = model.predict(input_img)
    60
    61        # (x1, y1, x2, y2) --> (c1, c2, w, h) (0.0, 1.0)
    62        boxes = np.array(boxes)[0]
    63        box_w  = boxes[:, 2] - boxes[:, 0]
    64        box_h = boxes[:, 3] - boxes[:, 1]
    65        boxes[:, 0] += box_w / 2
    66        boxes[:, 0] /= width
    67        boxes[:, 1] += box_h / 2
    68        boxes[:, 1] /= height
    69        boxes[:, 2] = box_w / width
    70        boxes[:, 3] = box_h / height
    71
    72        output = draw_bounding_boxes(orig_image,
    73                boxes,
    74                labels[0],
    75                VOC_CLASS_NAMES[1:],
    76                scores[0])
    77
    78        cv2.imshow('Faster RCNN Demo', output)
    79
    80        if cv2.waitKey(1) & 0xFF == ord('q'):
    81            break
    82
    83    cap.release()
    84    cv2.destroyAllWindows()

def frcnn_inference_demo(): View Source

    15def frcnn_inference_demo():
    16
    17    video = input(f"Please input the OpenCV capture device (e.g. 0, 1, 2): ")
    18
    19    while not video.isdigit():
    20        video = input(f"Please input the OpenCV capture device (e.g. 0, 1, 2): ")
    21
    22    # Capture from camera
    23    cap = cv2.VideoCapture(int(video))
    24    #cap.set(3, 1920)
    25    #cap.set(4, 1080)
    26
    27    # Download the model first if not exists
    28    # Check what_model_list for all available models
    29    index = 8
    30    if not os.path.isfile(os.path.join(WHAT_MODEL_PATH, what_model_list[index][WHAT_MODEL_FILE_INDEX])):
    31        get_file(what_model_list[index][WHAT_MODEL_FILE_INDEX],
    32                 WHAT_MODEL_PATH,
    33                 what_model_list[index][WHAT_MODEL_URL_INDEX],
    34                 what_model_list[index][WHAT_MODEL_HASH_INDEX])
    35
    36    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    37
    38    model = FasterRCNN(device=device)
    39
    40    model.load(os.path.join(WHAT_MODEL_PATH, what_model_list[index][WHAT_MODEL_FILE_INDEX]), map_location=device)
    41
    42    while True:
    43        _, orig_image = cap.read()
    44        if orig_image is None:
    45            continue
    46
    47        # Image preprocessing
    48        input_img = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
    49        height, width, _ = input_img.shape
    50
    51        # NHWC -> NCHW
    52        input_img = np.array(input_img).transpose((2, 0, 1))
    53        input_img = torch.from_numpy(input_img)[None]
    54
    55        # img = read_image('notebooks/demo.jpg', format='NCHW')
    56        # input_img = torch.from_numpy(img)[None]
    57        # RGB --> BGR
    58        # img = img.transpose((1, 2, 0)).astype(np.uint8)
    59
    60        inputs, boxes, labels, scores = model.predict(input_img)
    61
    62        # (x1, y1, x2, y2) --> (c1, c2, w, h) (0.0, 1.0)
    63        boxes = np.array(boxes)[0]
    64        box_w  = boxes[:, 2] - boxes[:, 0]
    65        box_h = boxes[:, 3] - boxes[:, 1]
    66        boxes[:, 0] += box_w / 2
    67        boxes[:, 0] /= width
    68        boxes[:, 1] += box_h / 2
    69        boxes[:, 1] /= height
    70        boxes[:, 2] = box_w / width
    71        boxes[:, 3] = box_h / height
    72
    73        output = draw_bounding_boxes(orig_image,
    74                boxes,
    75                labels[0],
    76                VOC_CLASS_NAMES[1:],
    77                scores[0])
    78
    79        cv2.imshow('Faster RCNN Demo', output)
    80
    81        if cv2.waitKey(1) & 0xFF == ord('q'):
    82            break
    83
    84    cap.release()
    85    cv2.destroyAllWindows()