8. OpenCV7

1. Object Detection(객체 탐지)

• 컴퓨터 비전과 이미지 처리와 관련된 컴퓨터 기술로써, 디지털 이미지와 비디오로 특정한 계열의 시맨틱 객체 인스턴스를 감지하는 일
• 얼굴 검출, 보행자 검출 등이 포함

 

2. 컴퓨터 비전의 Task 비교

• Image Classification: 이미지에 있는 개체 범주 목록 생성
• Single-Object Localization: 이미지에 있는 객체 범주의 한 인스턴스의 위치와 배율을 나타내는 Bounding Box를 생성
• Object Detection: 각 개체 범주의 모든 인스턴스의 위치와 배율을 나타내는 경계 상자와 함께 이미지에 있는 개체 목록을 생성
• 참조: https://oniss.tistory.com/39

 

3. Object Detection 실습 (util.py 넣기)

util.py

0.00MB

import os
import cv2
import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')
from collections import defaultdict
from ipywidgets import interact
from torch.utils.data import DataLoader
from torchvision import models, transforms
from torchvision.utils import make_grid
from torchvision.transforms import functional as F
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from util import CLASS_NAME_TO_ID, visualize, save_model
import util
from torchvision.ops import nms

data_dir = "./DataSet/"
data_df = pd.read_csv(os.path.join(data_dir, "df.csv"))
data_df

image_files = [fn for fn in os.listdir('./DataSet/train/') if fn.endswith('jpg')]
image_file = image_files[0]
image_file

image_path = os.path.join("./DataSet/train/", image_file)
image_path

image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image.shape
plt.imshow(image)

meta_data = data_df[data_df['ImageID'] == image_file.split(".jpg")[0]]
meta_data

cate_names = meta_data["LabelName"].values
cate_names

bboxes = meta_data[["XMin", "XMax", "YMin", "YMax"]].values
bboxes

class_ids = [CLASS_NAME_TO_ID[cate_name] for cate_name in cate_names]
class_ids

unnorm_bboxes = bboxes.copy()
# 'XMin', 'XMax', 'YMin', 'YMax' unnorm_bboxes
unnorm_bboxes

# XMin, XMax, YMin, YMax -> XMin, YMin, XMax, YMax
unnorm_bboxes[:, [1, 2]] = unnorm_bboxes[:, [2, 1]]
unnorm_bboxes

# XMin, XMax, YMin, YMax -> XMin, YMin, W, H
#W: (XMax - XMin), H: (YMax YMin)
unnorm_bboxes[:, 2:4] -= unnorm_bboxes[:, 0:2]
unnorm_bboxes

# XMin, YMin, W, H - X_Cen, Y_Cen, W, H
# X_Cen: (XMin + (W/2)), Y_Cen: (YMin + (H/2)) 
unnorm_bboxes[:, 0:2] += (unnorm_bboxes[:, 2:4]/2)
unnorm_bboxes

img_H, img_W, _ = image.shape
img_H, img_W

unnorm_bboxes[:, [0, 2]] *= img_W
unnorm_bboxes[:, [1, 3]] *= img_H
unnorm_bboxes

canvas = visualize(image, unnorm_bboxes, class_ids)
plt.figure(figsize=(6, 6))
plt.imshow(canvas)
plt.show()

@interact(index=(0, len(image_files)-1))
def show_sample(index=0):
    image_file = image_files[index]
    image_path = os.path.join('./DataSet/train/', image_file)
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image_id = image_file.split('.')[0]
    meta_data = data_df[data_df['ImageID'] == image_id]
    cate_names = meta_data['LabelName'].values
    bboxes = meta_data[['XMin', 'XMax', 'YMin', 'YMax']].values
    img_H, img_W, _ = image.shape
    class_ids = [CLASS_NAME_TO_ID[cate_name] for cate_name in cate_names]
    unnorm_bboxes = bboxes.copy()
    unnorm_bboxes[:, [1, 2]] = unnorm_bboxes[:, [2, 1]]
    unnorm_bboxes[:, 2:4] -= unnorm_bboxes[:, 0:2]
    unnorm_bboxes[:, 0:2] += (unnorm_bboxes[:, 2:4]/2)
    unnorm_bboxes[:, [0, 2]] *= img_W
    unnorm_bboxes[:, [1, 3]] *= img_H
    canvas = visualize(image, unnorm_bboxes, class_ids)
    plt.figure(figsize=(6, 6))
    plt.imshow(canvas)
    plt.show()

# 데이터셋 객체를 생성 후 인덱스를 입력하면 아래와 같은 데이터를 반환
# 이미지, Label, 파일이름
# 예) dataset = Dectection_dataset(..) / dataset[0] - > 이미지, 레이블, 파일이름

class Detection_dataset():
    def __init__(self, data_dir, phase, transformer=None):
        self.data_dir = data_dir
        self.phase = phase
        self.data_df = pd.read_csv(os.path.join(self.data_dir, 'df.csv'))
        self.image_files = [fn for fn in os.listdir(os.path.join(self.data_dir, phase)) if fn.endswith('jpg')]
        self.transformer = transformer
        
    def __len__(self):
        return len(self.image_files)
        
    def __getitem__(self, index):
        filename, image = self.get_image(index)
        bboxes, class_ids = self.get_label(filename)
        img_H, img_W, _ = image.shape
        if self.transformer:
            image = self.transformer(image)
            _, img_H, img_W = image.shape
        bboxes[:, [0, 2]] *= img_W
        bboxes[:, [1, 3]] *= img_H
        target = {}
        target['boxes'] = torch.Tensor(bboxes).float()
        target['labels'] = torch.Tensor(class_ids).long()
        return image, target, filename
        
    def get_image(self, index):
        filename = self.image_files[index]
        image_path = os.path.join(self.data_dir, self.phase, filename)
        image = cv2.imread(image_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        return filename, image
        
    def get_label(self, filename):
        image_id = filename.split('.')[0]
        meta_data = data_df[data_df['ImageID'] == image_id]
        cate_names = meta_data['LabelName'].values
        class_ids = [CLASS_NAME_TO_ID[cate_name] for cate_name in cate_names]
        bboxes = meta_data[['XMin', 'XMax', 'YMin', 'YMax']].values
        bboxes[:, [1, 2]] = bboxes[:, [2, 1]]
        return bboxes, class_ids

data_dir = "./DataSet/"
dataset = Detection_dataset(data_dir=data_dir, phase="train", transformer=None)

len(dataset)

dataset[0]

# 이미지의 shape는 (높이, 너비, 채널)로 표시
index=20
image, target, filename = dataset[index]
target, filename

boxes = target['boxes'].numpy()
class_ids = target['labels'].numpy()
n_obj = boxes.shape[0]
bboxes = np.zeros(shape=(n_obj, 4), dtype=np.float32)
bboxes[:, 0:2] = (boxes[:,0:2] + boxes[:, 2:4]) /2
bboxes[:, 2:4] = (boxes[:,2:4] - boxes[:, 0:2])
bboxes

canvas = util.visualize(image, bboxes, class_ids)

plt.figure(figsize=(6,6))
plt.imshow(canvas)
plt.show()

# dataset 객체에 index를 넣어 바운딩 박스를 표현하는 interact를 작성

@interact(index=(0, len(image_files) -1))
def show_sample(index=0):
    image, target, filename = dataset[index]

    boxes = target["boxes"].numpy()
    class_ids = target["labels"].numpy()
    
    n_obj = boxes.shape[0]
    bboxes = np.zeros(shape=(n_obj, 4), dtype=np.float32)
    bboxes[:, 0:2] = (boxes[:, 0:2] + boxes[:, 2:4]) / 2
    bboxes[:, 2:4] = boxes[:, 2:4] - boxes[:, 0:2]
    
    canvas = visualize(image, bboxes, class_ids)
    plt.figure(figsize=(6, 6))
    plt.imshow(canvas)

IMAGE_SIZE = 448

transformer = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize(size=(IMAGE_SIZE, IMAGE_SIZE)),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

data_dir = "./DataSet/"
transformed_dataset = Detection_dataset(data_dir=data_dir, phase="train", transformer=transformer)

image, target, filename = transformed_dataset[20]
image.shape

# make_grid(): 이미지 또는 그리드 형태의 데이터를 시각화 하기 위해 사용되는 함수
# 이미지를 그리드로 배열, 시각화를 일관성 있게 유지

np.image = make_grid(image, normalize=True).cpu().permute(1, 2, 0).numpy()
np.image.shape

# DataLoader(데이터셋, 배치사이즈, 셔플, collate_fn=collate_fn)
# collate_fn: 파이토치에서 데이터 로더에서 사용하는 함수. 배치 단위로 데이터를 나눌 때 사용
# 데이터 로더가 배치로 나눌 때 어떻게 처리할 지 정의함

def collate_fn(batch):
    image_list = []
    target_list = []
    filename_list = []

    for img, target, filename in batch:
        image_list.append(img)
        target_list.append(target)
        filename_list.append(filename)
        
    return image_list, target_list, filename_list

BATCH_SIZE = 8
trainset = Detection_dataset(data_dir=data_dir, phase="train", transformer=transformer)
trainloader = DataLoader(trainset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)

for index, batch in enumerate(trainloader):
    images = batch[0]
    targets = batch[1]
    filenames = batch[2]

    if index == 0:
        break

print(targets)

def build_dataloader(data_dir, batch_size=4, image_size=448):
    transformer = transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize(size=(image_size, image_size)),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # 안에 들어있는 숫자 수치는 ImageNet 데이터셋의 통계치
    ])
    dataloaders={}
    train_dataset = Detection_dataset(data_dir=data_dir, phase="train", transformer=transformer)
    dataloaders["train"] = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)

    val_dataset = Detection_dataset(data_dir=data_dir, phase="val", transformer=transformer)
    dataloaders["val"] = DataLoader(val_dataset, batch_size=1, shuffle=False, collate_fn=collate_fn)

    return dataloaders

data_dir = "./DataSet/"
dloaders = build_dataloader(data_dir)

for phase in ["train", "val"]:
    for index, batch, in enumerate(dloaders[phase]):
        images = batch[0]
        targets = batch[1]
        filenames = batch[2]
        
        print(targets)
        
        if index == 0:
            break

4. Two-Stage 모델

• 탐색 영역을 찾는 Region Proposal과 해당 영역을 분류하는 Detection 두 가자 과정이 순차적으로 수 행되는 방법
• 위치를 찾는 문제(Localization)
  • 하나의 이미지 안에서 물체가 있을법한 위치를 찾아 나열하는 과정에 대한 정보를 제안
• 분류 문제(Classification)
  • 각각의 위치에 대한 class를 분류
  • 이미지 내의 사물이 존재하는 bounding box를 예측하는 regression을 사용

4-1. R-CNN

• Selectve Search를 이용해 2,000개의 ROI를 추출
• 각 ROI에 대하여 동일한 크기의 입력 이미지로 변경
• 이미지를 CNN에 넣어서 벡터 이미지를 추출
• 해당 feature를 SVM에 넣어 class 분류 결과를 얻음

입력 이미지에 대해 CPU기반의 Selectve Search를 진행하므로 많은 시간이 소요

 

4-2. Fast R-CNN

• Selectve Search를 이용해 2,000개의 ROI를 추출
• 각 ROI에 대하여 동일한 크기의 입력 이미지로 변경
• 이미지를 CNN에 넣어서 벡터 이미지를 추출
• 해당 feature를 SVM에 넣어 class 분류 결과를 얻음

입력 이미지에 대해 CPU기반의 Selectve Search를 진행하므로 많은 시간이 소요

 

4-3. Faster R-CNN

• 속도가 느린 Region Proposal 작업을 GPU에서 수행함
• RPN(Region Proposal Networks) 적용
  • 슬라이딩 윈도우를 거쳐 각 위치에 대해 Regression과 classfication을 수행

 

5. One-Stage 모델

• Resion Proposal과 dectection이 한 번에 수행
• YOLO(You only Look Once): 2015년 제안된 객체 검출 모델로 이미지 전체를 단일 그리드로 나누고, 각 그리드 셀마다 여러 개의 바운딩 박스와 클래스를 예측하는 방식

 

6. Faster R-CNN (Resnet50)

model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model

def build_model(num_classes):
    model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model

NUM_CLASSES = 2
model = build_model(num_classes = NUM_CLASSES)
model

phase = 'train'
model.train()

for index, batch in enumerate(dloaders[phase]): 
    images = batch[0]
    targets = batch[1]
    filenames = batch[2]
    
    images = list(image for image in images) 
    targets = [{k: v for k, v in t.items()} for t in targets] 
    loss = model(images, targets)
    
    if index == 0:
        break

phase = 'train'
model.train()

for index, batch in enumerate(dloaders[phase]): 
    images = batch[0]
    targets = batch[1]
    filenames = batch[2]
    
    images = list(image for image in images) 
    targets = [{k: v for k, v in t.items()} for t in targets] 
    loss = model(images, targets)
    
    if index == 0:
        break

# loss_classifier: 객체 탐지 모델에서 분류기 손실 함수. 객체의 종류를 예측하는 사용
# loss_box_reg: 객체 위치를 예측하는 박스 회귀 모델의 손실 함수. 예측된 경계 상자의 위치와 실제 객체의 위치 사이의 차이를 줄이기 위해 사용
# loss_objectness: 객체 탐지 모델에서 사용되는 객체 존재 여부를 예측하는데 사용되는 손실함수. 각 경계 상자에 대해 해당 상자에 객체가 존재하는지 여부를 예측하고 실제와 비교하여 학습
# loss_rpn_box_reg: RPN의 박스 회귀 손실 함수. 객체 후보 영역을 제안하고 이 후보 영역의 경계 상자를 조정하기 위해 사용
loss

def train_one_epoch(dataloadrs, model, optimizer, device):
    train_loss = defaultdict(float)
    val_loss = defaultdict(float)
    model.train()
    
    for phase in ['train', 'val']:
        for index, batch in enumerate(dloaders[phase]):
            images = batch[0]
            targets = batch[1]
            filenames = batch[2]
            
            images = list(image for image in images)
            targets = [{k: v for k, v in t.items()} for t in targets]
            
            with torch.set_grad_enabled(phase == 'train'):
                loss = model(images, targets)
            total_loss = sum(each_loss for each_loss in loss.values())
            
            if phase == 'train':
                optimizer.zero_grad()
                total_loss.backward()
                optimizer.step()
                
                if (index > 0) and (index % VERBOSE_FREQ) == 0:
                    text = f"{index}/{len(dataloaders[phase])} - "
                    for k, v in loss.items():
                        text += f'{k}: {v.item():.4f} '
                    print(text)
                
                for k, v in loss.items():
                    train_loss[k] += v.item()
                train_loss['total_loss'] += total_loss.item()
            else:
                for k, v in loss.items():
                    val_loss[k] += v.item()
                val_loss['total_loss'] += total_loss.item()
    
    for k in train_loss.keys():
        train_loss[k] /= len(dataloaders['train'])
        val_loss[k] /= len(dataloaders['val'])
    return train_loss, val_loss

data_dir = './DataSet/'
is_cuda = False

NUM_CLASSES = 2 
IMAGE_SIZE = 448
BATCH_SIZE = 8
VERBOSE_FREQ = 100
DEVICE = torch.device('cuda' if torch.cuda.is_available and is_cuda else 'cpu')

dataloaders = build_dataloader(data_dir=data_dir, batch_size=BATCH_SIZE, image_size=IMAGE_SIZE)
model = build_model(num_classes=NUM_CLASSES)
model = model.to(DEVICE)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

num_epochs = 1

train_losses = []
val_losses = []

for epoch in range(num_epochs):
    train_loss, val_loss = train_one_epoch(dataloaders, model, optimizer, DEVICE) 
    train_losses.append(train_loss)
    val_losses.append(val_loss) 
    
    print(f"epoch:{epoch+1}/{num_epochs} - Train Loss: {train_loss['total_loss']:.4f}, Val Loss: {val_loss['total_loss']:.4f}")
    
    if (epoch+1) % 10 == 0:
        save_model(model.stat_dict(), f'model_{epoch+1}.pth')

tr_loss_classifier = []
tr_loss_box_reg = []
tr_loss_objectness = []
tr_loss_rpn_box_reg = []
tr_loss_total = []
for tr_loss in train_losses:
    tr_loss_classifier.append(tr_loss['loss_classifier'])
    tr_loss_box_reg.append(tr_loss['loss_box_reg'])
    tr_loss_objectness.append(tr_loss['loss_objectness'])
    tr_loss_rpn_box_reg.append(tr_loss['loss_rpn_box_reg'])
    tr_loss_total.append(tr_loss['total_loss'])
val_loss_classifier = []
val_loss_box_reg = []
val_loss_objectness = []
val_loss_rpn_box_reg = []
val_loss_total = []
for vl_loss in val_losses:
    val_loss_classifier.append(vl_loss['loss_classifier'])
    val_loss_box_reg.append(vl_loss['loss_box_reg'])
    val_loss_objectness.append(vl_loss['loss_objectness'])
    val_loss_rpn_box_reg.append(vl_loss['loss_rpn_box_reg'])
    val_loss_total.append(vl_loss['total_loss'])
    
plt.figure(figsize=(8, 4))
plt.plot(tr_loss_total, label="train_total_loss")
plt.plot(tr_loss_classifier, label="train_loss_classifier")
plt.plot(tr_loss_box_reg,  label="train_loss_box_reg")
plt.plot(tr_loss_objectness, label="train_loss_objectness")
plt.plot(tr_loss_rpn_box_reg,  label="train_loss_rpn_box_reg")
plt.plot(val_loss_total, label="train_total_loss")
plt.plot(val_loss_classifier, label="val_loss_classifier")
plt.plot(val_loss_box_reg,  label="val_loss_box_reg")
plt.plot(val_loss_objectness, label="val_loss_objectness")
plt.plot(val_loss_rpn_box_reg,  label="val_loss_rpn_box_reg")
plt.xlabel("epoch")
plt.ylabel("loss")
plt.grid("on")
plt.legend(loc='upper right')
plt.tight_layout()

def load_model(ckpt_path, num_classes, device):
    checkpoint = torch.load(ckpt_path, map_location=device)
    model = build_model(num_classes=num_classes)
    model.load_state_dict(checkpoint)
    model = model.to(device)
    model.eval()
    return model

model = load_model(ckpt_path='./DataSet/model_40.pth', num_classes=NUM_CLASSES, device=DEVICE)
model

 

7. Faster R-CNN (Resnet50)

• 객체 탐지와 같은 작업에서 사용되는 개념
• 객체 탐지 모델은 입력 이미지에서 객체의 위치를 찾아내는 작업을 수행 -> 모델은 주어진 이미지 내에 서 다양한 위치에 대해 객체가 존재하는지 예측하고 각 객체에 대한 바운딩 박스와 해당 객체에 대한 신뢰 도(Confidence Score)를 출력
• Confidence Threshold: 신뢰도를 조절하는 기준 값
  • 예) Confidence Threshold를 0.6으로 설정하면 모델은 신뢰도가 0.6이상인 객체만을 선택하게 됨
• Confidence Threshold를 적적하게 설정해야 객체 탐지의 정확도를 높일 수 있음, 너무 낮은 Confidence Threshold를 설정하면 신뢰성이 낮은 결과를 포함할 수 있고, 너무 높은 Confidence Threshold를 설정하면 신뢰성이 높은 객체조차 누락

 

8. Non-Maximum Suppression(NMS)

• 중복된 결과를 제거하여 정확하고 겹치지 않는 객체를 식별하는데 사용
• NMS가 작동되는 순서
  1. 객체 탐지 모델 실행
  2. 이미지를 입력받아 바운딩 박스와 신뢰도를 출력
  3. 바운딩 박스 필터링(겹치는 바운딩 박스들 중에서 가장 확실한 바운딩 박스(스코어 값이 가장 높음) 바운딩 박스만 남기고 나머지 겹치는 바운딩 박스는 제거 (IoU지표를 사용)

 

9. IoU(Intersection over Union)

• 객체 탐지나 세그멘테이션과 같은 컴퓨터 비전에서 모델이 예측한 결과와 실제 라벨 사이의 정확도를 측정하는 지표
• 바운딩 박스나 세그멘테이션 마스크가 얼마나 겹치는지를 측정하여 예측 결과의 정확성을 평가하는데 사용
• 0과 1 사이의 값으로 나타내며, 1에 가까울수록 예측 결과가 정확하고 겹치는 영역이 많다는 것을 의미
• 계산 방법
  1. 영역 A와 영역 B의 겹치는 영역을 계산(공통 부분을 계산) -> 교집합 계산(얼마나 겹쳐져 있는지)
  2. 합집합 계산(두 영역의 전체 크기)
  3. 교집합을 합집합으로 나눔 -> 교집합 / 합집합 = (IoU 계산)

def postprocess(prediction, conf_thres=0.3, IoU_threshold=0.3):
    pred_box = prediction['boxes'].cpu().detach().numpy()
    pred_label = prediction['labels'].cpu().detach().numpy()
    pred_conf = prediction['scores'].cpu().detach().numpy()

    valid_index = pred_conf > conf_thres 
    pred_box = pred_box[valid_index]
    pred_label = pred_label[valid_index]
    pred_conf = pred_conf[valid_index]

    valid_index = nms(torch.tensor(pred_box.astype(np.float32)), torch.tensor(pred_conf), IoU_threshold)
    pred_box = pred_box[valid_index.numpy()]
    pred_label = pred_label[valid_index.numpy()]
    pred_conf = pred_conf[valid_index.numpy()]
    
    return np.concatenate((pred_box, pred_conf[:, np.newaxis], pred_label[:, np.newaxis]), axis=1)

pred_images = []
pred_labels = []

for index, (images, _, filenames) in enumerate(dataloaders['val']): 
    images = list(image.to(DEVICE) for image in images)
    filename = filenames[0]

    image = make_grid(images[0].cpu().detach(), normalize=True).permute(1, 2, 0).numpy()
    image = (image * 255).astype(np.uint8)

    with torch.no_grad():
        prediction = model(images)

    prediction = postprocess(prediction[0])
    
    prediction[:, 2].clip(min=0, max=image.shape[1])
    prediction[:, 3].clip(min=0, max=image.shape[0])

    xc = (prediction[:,0] + prediction[:, 2]) / 2
    yc = (prediction[:,1] + prediction[:, 3]) / 2
    w = prediction[:, 2] - prediction[:, 0]
    h = prediction[:, 3] - prediction[:, 1]
    cls_id = prediction[:, 5]
    prediction_yolo = np.stack([xc, yc, w, h, cls_id], axis=1)

    pred_images.append(image)
    pred_labels.append(prediction_yolo)
    if index==2:
        break
        
pred_labels

@interact(index=(0, len(pred_images)-1))
def show_result(index=0):
    result = visualize(pred_images[index], pred_labels[index][:, 0:4], pred_labels[index][:, 4])
    plt.figure(figsize=(6,6))
    plt.imshow(result)
    plt.show()

video_path = './sample_video.mp4'

@torch.no_grad()
def model_predict(image, model):
    tensor_image = transformer(image)
    tensor_image = tensor_image.to(DEVICE)
    prediction = model([tensor_image])
    return prediction

video = cv2.VideoCapture(video_path)
while(video.isOpened()):
    ret, frame = video.read()
    if ret:
        ori_h, ori_w = frame.shape[:2]
        image = cv2.resize(frame, dsize=(IMAGE_SIZE, IMAGE_SIZE))
        prediction = model_predict(image, model)
        prediction = postprocess(prediction[0])
        prediction[:, [0,2]] *= (ori_w / IMAGE_SIZE)
        prediction[:, [1,3]] *= (ori_h / IMAGE_SIZE)
        prediction[:, 2].clip(min=0, max=image.shape[1])
        prediction[:, 3].clip(min=0, max=image.shape[0])
        
        xc = (prediction[:,0] + prediction[:, 2]) / 2
        yc = (prediction[:,1] + prediction[:, 3]) / 2
        w = prediction[:,2] - prediction[:, 0]
        h = prediction[:,3] - prediction[:, 1]
        cls_id = prediction[:, 5]
        prediction_yolo = np.stack([xc, yc, w, h, cls_id], axis=1)
        
        canvas = visualize(frame, prediction_yolo[:, 0:4], prediction_yolo[:, 4])
        cv2.imshow('camera', canvas)
        
        key = cv2.waitKey(1)
        if key == 27:
            break
        if key == ord('s'):
            cv2.waitKey()
video.release()