제로베이스 데이터취업스쿨 DAY86-107 자동차차종분류 YOLOv8x 어린이보호구역

# Object Detecion 
import cv2
from ultralytics import YOLO

#plots
import matplotlib.pyplot as plt
import seaborn as sns

#basics
import pandas as pd
import numpy as np
import os
from google.colab import drive
from tqdm import tqdm

# Display image and videos
import IPython
from IPython.display import Video, display
%matplotlib inline

import warnings
warnings.filterwarnings('ignore')

drive.mount('/content/drive')

model = YOLO('yolov8x.pt')
dict_classes = model.model.names

# 영상 크기 조절
def risize_frame(frame, scale_percent):
    """Function to resize an image in a percent scale"""
    width = int(frame.shape[1] * scale_percent / 100)
    height = int(frame.shape[0] * scale_percent / 100)
    dim = (width, height)

    # resize image
    resized = cv2.resize(frame, dim, interpolation=cv2.INTER_AREA)
    return resized
    
path = '/content/drive/MyDrive/딥러닝_프로젝트/2020_11_06_12_20_ride_cycle_sun_A_5.mp4'

verbose = False

# 사진 크기 비율(원본대비)
scale_percent = 100

video = cv2.VideoCapture(path)
class_IDS = [2, 3, 5, 7]

start_point = (5, 600)
end_point = (1589, 313)
offset = 5.5 

# 이전에 탐지된 차량들의 중심과 카운트를 저장하는 딕셔너리
prev_vehicles = {}

# 중복 인식 방지를 위한 시간 기반 필터링을 위한 시간 임계값 (단위: 프레임 수)
time_threshold = 0.5

# 각 객체들의 카운터
veiculos_contador = dict.fromkeys(class_IDS, 0)

# 원본 비디오의 크기
height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
fps = video.get(cv2.CAP_PROP_FPS)

# 비디오 크기 변경 (성능 향상을 위해)
if scale_percent != 100:
    width = int(width * scale_percent / 100)
    height = int(height * scale_percent / 100)

# 비디오 출력 설정
video_name = 'result.mp4'
output_path = "rep_" + video_name
tmp_output_path = "tmp_" + output_path
VIDEO_CODEC = "MP4V"

output_video = cv2.VideoWriter(tmp_output_path, 
                               cv2.VideoWriter_fourcc(*VIDEO_CODEC), 
                               fps, (width, height))

# 탐지 실행
for i in tqdm(range(int(video.get(cv2.CAP_PROP_FRAME_COUNT)))):
    # 비디오로부터 프레임 읽기
    _, frame = video.read()
    
    # 프레임 크기 변경
    #frame = risize_frame(frame, scale_percent)
    
    if verbose:
        print('크기 조정된 프레임 크기: ', (frame.shape[1], frame.shape[0]))

    # 탐지 실행
    y_hat = model.predict(frame, conf=0.5, classes=class_IDS, device=0, verbose=False)
    
    # 현재 프레임에서 탐지된 객체들의 바운딩 박스, 신뢰도 및 클래스 가져오기
    boxes = y_hat[0].boxes.xyxy.cpu().numpy()
    conf = y_hat[0].boxes.conf.cpu().numpy()
    classes = y_hat[0].boxes.cls.cpu().numpy() 
    
    # 중복된 탐지 및 선을 넘어가지 않은 차량들 필터링
    filtered_boxes = []
    filtered_conf = []
    filtered_classes = []
    
    for j in range(len(boxes)):
        # 각 차량의 좌표 얻기
        xmin, ymin, xmax, ymax = boxes[j].astype('int')
        
        # 바운딩 박스의 중심 계산
        center_x, center_y = int((xmax + xmin) / 2), int((ymax + ymin) / 2)
        
        # 차량을 고유 식별자로 생성 (중심 좌표를 기반으로)
        vehicle_id = f'{center_x}_{center_y}'
        
        # 이전 프레임에서 이 차량이 감지되지 않았는지 확인
        if vehicle_id not in prev_vehicles:
            filtered_boxes.append(boxes[j])
            filtered_conf.append(conf[j])
            filtered_classes.append(classes[j])
            # 차량의 중심 좌표와 프레임 카운트를 딕셔너리에 저장
            prev_vehicles[vehicle_id] = (center_x, center_y, i)
        
        # 선을 넘어가지 않은 차량인지 확인
        m = (end_point[1] - start_point[1]) / (end_point[0] - start_point[0])
        b = start_point[1] - m * start_point[0]
        if center_y < (m * center_x + b + offset) and center_y > (m * center_x + b - offset):
            
            # 이전 프레임에서 인식된 차량과 겹치지 않으면서 시간 임계값 이내에 있는 경우에만 카운트 증가
            if vehicle_id in prev_vehicles:
                prev_center_x, prev_center_y, prev_frame_count = prev_vehicles[vehicle_id]
                if abs(prev_frame_count - i) <= time_threshold:
                    veiculos_contador[classes[j]] += 1
    
    # 차량 종류별 수를 갱신
    xt = 40
    for category in veiculos_contador:
      xt += 30
      cv2.putText(img=frame, text=f'{dict_classes[category]}: {veiculos_contador[category]}', 
                  org=(30, xt), fontFace=cv2.FONT_HERSHEY_TRIPLEX, 
                  fontScale=1, color=(135, 206, 235), thickness=2)

    if verbose:
        print(veiculos_contador)
    
    # 변환된 프레임을 비디오에 저장
    output_video.write(frame)

# 비디오 저장 종료
output_video.release()

# 최종 결과 데이터프레임 생성
result_df = pd.DataFrame([veiculos_contador])
result_df.rename(columns=dict_classes, inplace=True)

# 결과 데이터프레임 출력
result_df

영상 출처 : https://aihub.or.kr/aihubdata/data/view.do?currMenu=115&topMenu=100&aihubDataSe=realm&dataSetSn=169

참고 코드 : https://www.kaggle.com/code/paulojunqueira/yolo-v8-vehicles-detecting-counting/notebook