Object Detection Model(YOLO+AR)

지난번에 배운 YOLO 복습도 할 겸

이번에는 영상에서 고양이를 탐지한 다음 고양이 머리쪽에 왕관을 합성해서 새로운 영상으로 저장해주는 코드를 만들어 보자

 

자료출처

고양이 영상: pexels.com에서 무료 다운로드

왕관그림: chatgpt한테 요청

haarcascade_frontalcatface_extended.xml: github에서 다운로드

 

최종코드

model = YOLO('yolov8n.pt') 

video_path = 'earthcat.mp4'
cap = cv2.VideoCapture(video_path)
out = cv2.VideoWriter('output_with_crown.mp4', cv2.VideoWriter_fourcc(*'mp4v'), cap.get(cv2.CAP_PROP_FPS), (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))))

cat_face = cv2.CascadeClassifier('haarcascade_frontalcatface_extended.xml')
crown = cv2.imread('crown.png', cv2.IMREAD_UNCHANGED)

while cap.isOpened():    
    ret, frame = cap.read()
    if not ret:
        break

    results = model(frame)

    for result in results[0].boxes:
        if int(result.cls) == 15:
            x1,y1,x2,y2=map(int,result.xyxy[0])
            cat_region=frame[y1:y2, x1:x2]
            gray_region=cv2.cvtColor(cat_region,cv2.COLOR_BGR2GRAY)
            faces=cat_face.detectMultiScale(gray_region, scaleFactor=1.1, minNeighbors=3)

        for (fx,fy,fw,fh) in faces:
            hx=x1+fx
            hy=y1+fy
            hw=fw
            hh=fh

            crown_resized = cv2.resize(crown, (hw, int(hh * 0.5)))
            crown_rgb = crown_resized[:, :, :3]  # RGB 채널
            crown_alpha = crown_resized[:, :, 3]  # Alpha 채널 (투명도)

            cx = hx
            cy = hy - int(hh * 0.5)
            roi = frame[cy:cy + crown_resized.shape[0], cx:cx + crown_resized.shape[1]]

            alpha = crown_alpha / 255.0
            alpha_inv = 1.0 - alpha

            for c in range(3):  
                roi[:, :, c] = (alpha * crown_rgb[:, :, c] + alpha_inv * roi[:, :, c])

            frame[cy:cy + crown_resized.shape[0], cx:cx + crown_resized.shape[1]] = roi

    out.write(frame)
    cv2.imshow('Crowned Cat', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
out.release()
cv2.destroyAllWindows()

 

우선 영상에서 yolo를 사용해 영상 내부의 객체를 탐지해보자

이전에 배웠던 코드에서 이미지만 비디오로 바꾸고 VideoCapture를 쓰면 된다.

이 때, cv2.VideoCapture(0)을 하면 (0번카메라, 주로 기본캠에서) 실시간 영상탐지로 동작하고, 0 대신 영상물의 경로를 입력하면 저장된 영상에서 객체탐지를 한다.

cv2.cvtColor(사진, 변환코드) 는 이후에도 자주 나오는데, 간단히 말해 OpenCV에서 색상처리변환 함수이다.

OpenCV는 이미지를 BGR로 읽는데, 다른 라이브러리에서 동작할 때는 RGB 또는 grayscale로 변환해야 할 때가 있다.

 

코드

from ultralytics import YOLO
import cv2
from matplotlib import pyplot as plt

 

model = YOLO('yolov8n.pt') 

video_path = 'earthcat.mp4'
cap = cv2.VideoCapture(video_path)

while True:
    ret, frame = cap.read()
    if not ret:
        break

    results = model(frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

img_with_boxes = result.plot()

plt.imshow(cv2.cvtColor(img_with_boxes, cv2.COLOR_BGR2RGB))
plt.axis('off')
plt.show()

이 코드도 동작은 하지만, 문제가 있다. 

영상은 2분짜리인데, 모든 프레임을 다 처리하니 시간이 너무 오래걸린다.

그래서 우선 프레임 몇 개를 스킵하고 일부만 객체탐지를 진행하면서, 코드가 잘 동작하는지 보고

나중에 스킵하는 부분을 없애도로 하자

 

스킵코드 추가

model = YOLO('yolov8n.pt') 

video_path = 'earthcat.mp4'
cap = cv2.VideoCapture(video_path)

frame_skip = 50 

while True:
    for _ in range(frame_skip - 1):  
        ret = cap.grab()  
        if not ret:
            break
    
    ret, frame = cap.read()
    if not ret:
        break


    results = model(frame)
  
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

img_with_boxes = results[0].plot()

plt.imshow(cv2.cvtColor(img_with_boxes, cv2.COLOR_BGR2RGB))
plt.axis('off')
plt.show()

 

frame_skip = 50 

while True:
    for _ in range(frame_skip - 1):  
        ret = cap.grab()  
        if not ret:
            break
    
    ret, frame = cap.read()
    if not ret:
        break

코드를 위에서부터 읽어보면 50프레임씩 스킵을 할 것이고,

50-1=49 -> for _ in range(49): 이하의 루프를 49회 반복하면서 cap.grab() 이 작동한다.

cap.grab()은 프레임을 실제로 처리하지 않고 다음 프레임으로 넘겨준다. 49번째 프레임까지 이렇게 처리되고 for_in 문이 끝난다.

이후 ret.frame=cap.read()가 50번째 프레임을 처리한다.

 

그러면 while True 안에 있는 모든 코드가 끝난 것이므로 다시 while True의 처음부분으로 돌아가서,

cap.grab이 다시 49회 반복된다. (51번 프레임부터 99번 프레임까지 처리없이 넘어감)

그리고 for_in 문이 끝났으므로 ret.frame=cap.read()가 작동하여 100번째 프레임은 cap.read()가 처리한다.

이 때, 1번프레임부터 다시 시작하지 않는 이유는 cap.grab()과 cap.read()가 프레임을 옮겨놓는 역할까지 하기 때문이다. 50번까지 처리가 끝났으면 프레임은 50번째에 계속 머물러있지 1번으로 돌아가지 않는다.

 

if not ret:

  break

이 구문은 cap.grab()과 cap.read() 이후에 나오는데, 만약 비디오가 끝나거나 뭔가 오류가 발생하면 ret(return이라는 뜻)이 반환되지 않으니 ret=False가 되어 반복문을 끝낸다.

 

 results[0].plot()

기본코드는 원래 results.plot()이었는데 이렇게 동작시켰더니 리스트 형식은 plot 매서드를 쓸 수 없다고 오류가 난다.

그래서 리스트의 내용물을 빼내기 위해 [0]을 넣어줬다.

이 모델은 매 프레임마다 단일프레임을 탐지하므로 result[0]만 존재하지만, 만약 여러 프레임을 동시에 처리하는 모델이라면 result[1]은 두번째 프레임의 탐지 결과가, result[2]에는 세번째 프레임의 탐지결과가 있다. 

 

Q 비디오니까 여러 프레임을 처리한거 아닌지?

A while True에서 한번에 한 프레임만 가져와서 모델에 넣고 있으므로 하나의 프레임에서만 객체탐지가 이루어지고 있다.

여기까지 돌려봤을 때 결과

탐지된 객체와 객체의 shape이 표시되고 맨 마지막에 고양이가 탐지된 사진이 나온다.

저 사진은 영상의 거의 끝부분에 나오는 고양이의 모습이다.

매번 results=model(frame) 에서 탐지된 결과가 덮어쓰기 되고 있어서, 최종적으로 가장 마지막 프레임에서 탐지된 결과만 나왔다.

그렇다면 프레임마다의 탐지 결과를 모두 볼수는 없을까?

 

모든 프레임의 결과 보는 코드

model = YOLO('yolov8n.pt') 

video_path = 'earthcat.mp4'


cap = cv2.VideoCapture(video_path)

frame_skip = 50  

cat=[]

while True:
    for _ in range(frame_skip - 1):  
        ret = cap.grab()  
        if not ret:
            break
    
    ret, frame = cap.read()
    if not ret:
        break

    results = model(frame)

    if any(int(box.cls) == 15 for box in results[0].boxes):
        img_with_boxes = results[0].plot()
        cat.append(img_with_boxes)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

for i, img_with_boxes in enumerate(cat):
    plt.figure()  
    plt.imshow(cv2.cvtColor(img_with_boxes, cv2.COLOR_BGR2RGB))
    plt.axis('off')
    plt.title(f'Cat Detection Frame {i+1}') 
    plt.show()

나중에 고양이에다만 왕관을 합성해야 하니, 지금부터 고양이로 탐지된 부분만을 모아서 따로 저장해보자

 

우선 고양이의 클래스가 yolov8n에서 얼마로 지정되었는지를 알아봐야 한다.

model = YOLO('yolov8n.pt')
print(model.names)

고양이는 15번 클래스이다.

 

cat=[]

    if any(int(box.cls) == 15 for box in results[0].boxes):
          img_with_boxes = results[0].plot()
          cat.append(img_with_boxes)

 

이렇게 하면 cat이라는 리스트 안에다가 append로 고양이가 탐지된 이미지 박스만을 모아둘 수 있다.

results.boxes에는 탐지한 객체의 클래스이름, 신뢰도 점수가 들어있다. 거기서 box.cls==15 즉 고양이인 것만 모으면 된다.

여기까지의 결과

 

고양이만 탐지한 박스들을 얻었다.

여기까지는 쉬웠다.

 

그런데.. 생각해보니 내가 합성할건 왕관이고 그러려면 고양이 자체가 아니라 고양이의 머리가 필요한데

탐지된 박스를 보면 고양이 머리가 오른쪽에도 있고 가운데도 있고 아래도 있고 왼쪽에도 있고 제각각이다.

즉 왕관을 어느 위치에 놓아야 할지 좌표를 정할수가 없다. -> 왕관을 제대로 합성할 수 없다.

(등에다 날개 합성했으면 더 쉬웠을까..?)

그래서

 

고양이 얼굴 탐지 모델(haarcascade_frontalcatface_extended.xml)로 얼굴을 탐지하고

yolo에서 탐지한 고양이의 좌표와, cascade에서 탐지한 고양이 얼굴 좌표를 이용해 왕관의 좌표를 설정해야겠다!

 

우선 좌표를 얻는 코드는

print(results[0].boxes.xyxy) #박스 좌상단모서리의 x좌표, y좌표, 박스 우하단모서리의 x좌표, y좌포
print(results[0].boxes.xywh) #박스의 정가운데의 x좌표, y좌표, 박스너비, 박스높이 

텐서형태로 반환됨

이제 최종코드에서 이전과 달리 새로 추가된 부분을 살펴보면

out = cv2.VideoWriter('output_with_crown.mp4', cv2.VideoWriter_fourcc(*'mp4v'), cap.get(cv2.CAP_PROP_FPS), (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))))

최종 결과물을 저장하는 코드

cat_face = cv2.CascadeClassifier('haarcascade_frontalcatface_extended.xml')

고양이 얼굴 탐지모델(github에서 다운로드 가능)

https://github.com/opencv/opencv/tree/master/data/haarcascades

crown = cv2.imread('crown.png', cv2.IMREAD_UNCHANGED)

왕관이미지를 가져오는 코드인데, 이미지의 배경을 투명하게 처리해야 자연스럽게 합성되며 그렇게 하기 위해 배경이 투명하게 처리된 png파일로 가져와야 한다. 만약 배경 투명화가 되지 않았으면 구글에 검색해서 무료로 1초만에 만들수 있다.

while cap.isOpened():   

이전에 while True: 문을 썼지만, cap.isOpened()로 바꿔서 cv2.VideoCapture의 객체인 cap이 정상적으로 비디오 파일을 열어서 isOpened()=True일 때만 코드가 동작하게 했다. 비디오가 끝나면 코드도 종료하게 만들기 위해서

    for result in results[0].boxes:
        if int(result.cls) == 15:
            x1,y1,x2,y2=map(int,result.xyxy[0])
            cat_region=frame[y1:y2, x1:x2]
            gray_region=cv2.cvtColor(cat_region,cv2.COLOR_BGR2GRAY)
            faces=cat_face.detectMultiScale(gray_region, scaleFactor=1.1, minNeighbors=3)

고양이로 분류된 바인딩박스에서

좌표를 구하고 x1,y1,x2,y2=map(int,result.xyxy[0])

고양이 영역으로 구분하고 이 때, OpenCV는 이미지를 행렬 형식으로 처리하기 때문에 y좌표(세로, 행)가 먼저 나왔다.

frame[y1:y2, x1:x2]  -> 프레임에서 이 영역만큼 슬라이싱 하라는 뜻

그리고 haarcascade의 detectmultiscale 흑백에서 작동하기 때문에 BGR2GRAY 처리를 해줬다. 

cat_face.detectMultiScale() 에서 괄호안의 인자들은 이미지 탐지 매개변수들이다.

최종적으로 faces안에는 탐지한 고양이 얼굴 바운딩 박스의 (x,y,w,h좌표) 정보가 리스트형태로 들어있다.

ex)  faces=[(50, 60, 100, 100), (150, 200, 80, 80),...]

 

코드의 흐름을 보면

yolo로 고양이 영역 탐지

-> 고양이 영역을 흑백으로 처리하고

-> 고양이 영역 안에서 haarcascade가 고양이 얼굴영역 탐지

        for (fx,fy,fw,fh) in faces:
            hx=x1+fx
            hy=y1+fy
            hw=fw
            hh=fh

이건 고양이 머리의 좌표를 전체 사진 안에서 지정하는 코드이다.

x1,y1,x2,y2는 yolo가 탐지한 고양이 바운딩박스의 좌표

fx,fy,fw,fh는 위 고양이 바운딩박스 내부에서 haarcascade가 탐지한 고양이 얼굴좌표(상대적인 좌표) 

hx,hy,hw,hh는? 전체 사진에서 절대적인 고양이 얼굴좌표

 

사실.. yolo 안쓰고 처음부터 haarcascade만 써서 머리만 탐지하면 훨씬 간편할것 같기도 하다.

다만 영상이 복잡해지면 고양이만 탐지하고 그 안에서 머리를 탐지하는 식으로 해야 더 나은 결과가 나올것이다.

 

            crown_resized = cv2.resize(crown, (hw, int(hh * 0.5)))
            crown_rgb = crown_resized[:, :, :3]  
            crown_alpha = crown_resized[:, :, 3]  

            cx = hx
            cy = hy - int(hh * 0.5)
            roi = frame[cy:cy + crown_resized.shape[0], cx:cx + crown_resized.shape[1]]

            alpha = crown_alpha / 255.0
            alpha_inv = 1.0 - alpha

            for c in range(3):  
                roi[:, :, c] = (alpha * crown_rgb[:, :, c] + alpha_inv * roi[:, :, c])

            frame[cy:cy + crown_resized.shape[0], cx:cx + crown_resized.shape[1]] = roi

 

이제 왕관을 처리해야 하는데 솔직히 이 부분이 전체 코드중에서도 제일 어렵고 헷갈리고 오류나고 힘들었다.

특히 가장 고비었던 부분만 얼른 훑고 지나가자

(솔직히 다 이해도 못했다. 근데 코드는 동작한다 how...?)

            crown_rgb = crown_resized[:, :, :3]  
            crown_alpha = crown_resized[:, :, 3]  

rgb는 색깔이 있는 rgb채널을 의미하고, alpha는 알파채널 즉 투명도 정보를 의미한다.(0=가장투명, 255=가장불투명)

채널는 rgb+alpha가 있고 rgb는 0부터 2까지에 할당되고 alpha는 3에 할당되어 0부터 3까지 총 4개이다.

 [:, : , :3] 에서 :3은 앞에서부터 3 이전까지 =0,1,2

 [:, :, 3] 은 그냥 3번 

늘 말하지만 슬라이싱 방식을 헷갈리면 안된다. 뒷부분 포함하지 않음!

 

그 밑의 cx, cy는 합성할 왕관의 x,y좌표인데

cy = hy - int(hh * 0.5)

이걸 보면 머리의 y좌표에서 머리높이의 절반만큼을 뺀 곳에 왕관을 위치시키고 있다.

Q 머리 위에 왕관을 올리려면 더해야 하지 않나?

A OpenCV는 좌상단의 좌표가 (0, 0)이고 아래로 갈수록 y값이 커진다.  그래서 왕관이 머리보다 위로 올라가려면 숫자가 줄어들어야한다.

 

roi는 원본frame에서 왕관이 들어갈 자리를 잘라낸 부분이다.

 

for c in range(3): 

이 코드는 R,G,B 셋에 대해 처리하겠다는 의미

 

    out.write(frame)

최종적으로 처리한프레임을 제일 위의 out=... 에 기록함 ->비디오로 저장됨

cap.release()
out.release()
cv2.destroyAllWindows()

비디오 처리가 끝나면 리소스를 해제하고 프로그램을 종료하게 만드는 코드들

안해주면 만약 실시간 탐지였으면 웹캠이 계속 돌아간다든지 메모리누수라든지 이런저런 문제들이 생길수있다

 

 

최종결과..! (영상에서 캡쳐함)

 

!!!!

모든 프레임에서 항상 성공한것은 아니지만

그래도 최근 만든 모델 중에는 가장 그럴듯한 결과물이 나와서 기쁘다

YOLO 배울때는 코드가 간단해서 쉽게 생각했는데 생각보다 쉽지 않다..

OpenCV에서 어떻게 동작할지를 이해하고 좌표를 생각해내는 부분이 특히나 어려웠다.

왕관 처리는 말할것도 없고..

 

 

 

what's next?

알고리즘 코드카타

신경안쓰고 있었는데 해보면 좋을 것 같다.