학생 때 C++ D 받았었는데, AI 문과생으로 코딩이 너무 재미있어졌어요

import torch
import cv2
import os
import numpy as np
from ultralytics import YOLO
from filterpy.kalman import KalmanFilter

def compute_iou(bbox1, bbox2):
    """
    두 개의 상자가 얼마나 겹치는지 계산하는 함수예요.
    겹치는 정도를 0에서 1 사이의 숫자로 알려줍니다.
    """
    x1, y1, x2, y2 = bbox1
    x3, y3, x4, y4 = bbox2

    x_left = max(x1, x3)
    y_top = max(y1, y3)
    x_right = min(x2, x4)
    y_bottom = min(y2, y4)

    if x_right < x_left or y_bottom < y_top:
        return 0.0

    intersection_area = (x_right - x_left) * (y_bottom - y_top)
    bbox1_area = (x2 - x1) * (y2 - y1)
    bbox2_area = (x4 - x3) * (y4 - y3)
    union_area = bbox1_area + bbox2_area - intersection_area

    iou = intersection_area / union_area
    return iou

def assign_id(bbox, prev_bboxes, threshold=0.5):
    """
    새로운 상자에 ID를 부여하는 함수예요.
    이전 상자들과 비교해서 가장 비슷한 상자의 ID를 가져옵니다.
    """
    if not prev_bboxes:
        return len(prev_bboxes)

    ious = [compute_iou(bbox, prev_bbox) for prev_bbox in prev_bboxes]
    max_iou = max(ious)

    if max_iou >= threshold:
        return ious.index(max_iou)
    else:
        return len(prev_bboxes)

def create_kalman_filter():
    """
    칼만 필터라는 것을 만드는 함수예요.
    칼만 필터는 상자의 움직임을 부드럽게 예측하는 데 사용됩니다.
    """
    kf = KalmanFilter(dim_x=4, dim_z=2)
    kf.F = np.array([[1, 0, 1, 0],
                     [0, 1, 0, 1],
                     [0, 0, 1, 0],
                     [0, 0, 0, 1]])
    kf.H = np.array([[1, 0, 0, 0],
                     [0, 1, 0, 0]])
    kf.P *= 1000
    kf.R *= 100
    kf.Q *= 0.01
    return kf

def update_kalman_filter(kf, bbox):
    """
    칼만 필터를 업데이트하는 함수예요.
    상자의 새로운 위치를 칼만 필터에 알려줍니다.
    """
    x, y, w, h = bbox
    center_x, center_y = x + w // 2, y + h // 2
    kf.predict()
    kf.update(np.array([center_x, center_y]))
    return kf.x[:2]

# 입력 비디오 경로와 출력 폴더 경로를 지정해 주세요.
video_path = "input_video.mp4"
highlight_path = "highlights"

# YOLOv8 모델을 불러옵니다. 이 모델은 사람을 찾는 데 사용될 거예요.
model = YOLO('yolov8s.pt')

# 비디오를 불러옵니다.
video = cv2.VideoCapture(video_path)

# 비디오 파일이 열렸는지 확인합니다.
if not video.isOpened():
    print("비디오 파일을 열 수 없습니다.")
    exit()

# 사람을 추적하기 위한 변수들을 초기화합니다.
person_detections = {}
prev_bboxes = []
kalman_filters = {}
frame_num = 0

while True:
    # 비디오에서 이미지를 한 장씩 읽어옵니다.
    ret, frame = video.read()
    if not ret:
        break

    # YOLOv8 모델을 사용하여 이미지에서 사람을 찾습니다.
    results = model(frame)

    # 찾은 사람들 중에서 정말 사람인 것만 골라냅니다.
    person_info = results[0].boxes.cls == 0

    # 검출된 사람 객체의 수를 출력합니다.
    print(f"검출된 사람 객체 수: {len(results[0].boxes.cls[person_info])}")

    # 사람 주변에 초록색 상자를 그립니다.
    curr_bboxes = []
    for box in results[0].boxes.xyxy[person_info]:
        x1, y1, x2, y2 = map(int, box.tolist())
        curr_bboxes.append((x1, y1, x2, y2))
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)

    # 칼만 필터를 사용하여 사람을 부드럽게 추적합니다.
    for bbox in curr_bboxes:
        person_id = assign_id(bbox, prev_bboxes)
        if person_id not in person_detections:
            person_detections[person_id] = []
            kalman_filters[person_id] = create_kalman_filter()
        person_detections[person_id].append((frame_num, bbox))
        prev_center = update_kalman_filter(kalman_filters[person_id], bbox)
        prev_center_x, prev_center_y = map(int, prev_center)
        cv2.circle(frame, (prev_center_x, prev_center_y), 5, (0, 0, 255), -1)

    prev_bboxes = curr_bboxes

    # 처리된 이미지를 화면에 보여줍니다.
    cv2.imshow('frame', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

    frame_num += 1

video.release()
cv2.destroyAllWindows()

# 하이라이트 영상을 저장할 폴더를 만듭니다.
os.makedirs(highlight_path, exist_ok=True)

# 각 사람별로 하이라이트 영상을 만듭니다.
for person_id, detections in person_detections.items():
    highlight_frames = []

    for frame_num, bbox in detections:
        x1, y1, x2, y2 = bbox

        video.set(cv2.CAP_PROP_POS_FRAMES, frame_num)
        ret, frame = video.read()

        if ret:
            # 사람 주변을 크게 잘라냅니다.
            margin = 50
            x1, y1, x2, y2 = x1 - margin, y1 - margin, x2 + margin, y2 + margin
            x1, y1 = max(0, x1), max(0, y1)
            x2, y2 = min(frame.shape[1], x2), min(frame.shape[0], y2)

            person_frame = frame[y1:y2, x1:x2]

            # 잘라낸 이미지를 줌인/줌아웃 효과를 주어 강조합니다.
            zoom_ratio = 1.5
            zoom_in_frame = cv2.resize(person_frame, None, fx=zoom_ratio, fy=zoom_ratio)
            zoom_out_frame = cv2.resize(zoom_in_frame, (person_frame.shape[1], person_frame.shape[0]))

            highlight_frames.append(zoom_out_frame)

    # 하이라이트 영상을 저장합니다.
    if highlight_frames:
        highlight_filename = os.path.join(highlight_path, f"person_{person_id}.mp4")
        height, width, _ = highlight_frames[0].shape
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        out = cv2.VideoWriter(highlight_filename, fourcc, 30.0, (width, height))
        for frame in highlight_frames:
            out.write(frame)
        out.release()
        print(f"사람 {person_id}의 하이라이트 영상이 {highlight_filename}에 저장되었습니다.")
    else:
        print(f"사람 {person_id}의 하이라이트 영상이 없습니다.")

video.release()

print("하이라이트 영상 생성이 완료되었습니다!")