[빅데이터분석] 딥러닝_16. 사진 속에서 검출 인터페이스 만들기

책 그림 8-17

사물검출이란 다수의 사물이 존재하는 상황에서 각 사물의 위치와 클래스를 찾는 기술

 

현업사례:  

 

1. 딥노이드 회사의 선배는 인천공항에서 위반된 물품을 사물검출하는 기술을 구현하고 유지보수를 하고 있습니다.

 2. 딥노이드 회사의 폐사진에서 폐결절 부분을 찾아내는 사물검출

3. 일본에서 도심에서 거리의 유동인구를 카운트 해서 인구가 많아지면 경고 메세지를 발동

 

책 그림 8-18

1. 원본 이미지가 신경망에 입력이 됩니다.

2. 입력이미지에서 후보영역(region proposals)를 여러개 추출합니다. 후보영역은 물체가 있을 가능성이 높은 부분을 뜻합니다.

3. 후보영역중에 물체로 보이는 가장 확률이 높은 사진을 추출

4. cnn 신경망에 입력해서 분류합니다. 

 

■ 성능을 올리는 방법 

1. 데이터 증강 :  데이터를 다양하게 변형해서 학습 데이터의 다양성을 높이는  방법을 사용합니다. 

 욜로는 자동으로 이미지를 증식 시키는 기능이 있습니다.

# 모델 학습 실행
model.train(
    data=data_yaml,  # 학습에 사용할 데이터 경로 (data.yaml 파일)
    epochs=100,  # 학습을 반복할 횟수
    imgsz=640,  # 640x640 해상도 크기 , 이미지 크기
    batch=16,  # 한번에 학습할 이미지의 갯수
    name='pineapple_apple_detector',  # 프로젝트 이름
    workers=8  # 데이터 로딩에 사용할 CPU 코어 수,
    augment= True   # 데이터 증강 자동 활성화
)

2.  더 큰 Yolo 모델 사용

# 1. 가장 가벼운 모델 :   model = YOLO('yolov8v.pt')

# 2. 중간 모델 :  model = YOLO('yolov8m.pt')

# 3. large 모델 :   model = YOLO('yolov8l.pt')

 

3. 배치 사이즈 조정 

 한번에 학습할 이미지의 수를  사람이 직접 조정하면서 정확도가 올라가는지 확인

model.train(
    data=data_yaml,  # 학습에 사용할 데이터 경로 (data.yaml 파일)
    epochs=100,  # 학습을 반복할 횟수
    imgsz=640,  # 640x640 해상도 크기 , 이미지 크기
    batch=16,  # 한번에 학습할 이미지의 갯수
    name='pineapple_apple_detector',  # 프로젝트 이름
    workers=8  # 데이터 로딩에 사용할 CPU 코어 수,
    augment= True   # 데이터 증강 자동 활성화
)

4. 에포크수 조정  
 
 에포크를 늘리면 학습 시간이 더 오래걸리지만 결과는 더 좋은 결과를 얻을 수 있습니다. 

model.train(
    data=data_yaml,  # 학습에 사용할 데이터 경로 (data.yaml 파일)
    epochs=100,  # 학습을 반복할 횟수

5. 전이학습  (가장 중요한 부분)

 사전 학습된 yolo 모델을 기반으로 추가 학습을 진행할 때, 이전에 학습된 가중치를 
 유지하면서 미세조정을 수행할 수 있습니다. 

model= YOLO('yolov8m.pt')  # 사전에 학습된 모델 
model.train( 
               data = data_yaml,
               epochs =50, 
               pretrained=True  # 사전에 학습된 가중치를 유지하면서 학습 
             )

 

 

 

 

 

 

---

※ 실

 

1. 필요한 모듈 설치

!pip install ultralytics

 

※ 압축파일 풀기 (학습할 파일 코랩&mydrive에 올리기)

!unzip /content/drive/MyDrive/yolo9/val.zip -d /content/drive/MyDrive/yolo9/

 

2. 사물검출 신경망 학습시 훈련 데이터와 테스트 데이터의 위치와 정답정보를 알려주는 메타 데이터

import yaml

# 데이터셋 경로 및 클래스 정보 설정
data_yaml = {
    'train': '/content/drive/MyDrive/yolo9/train',  # Train 데이터 경로
    'val': '/content/drive/MyDrive/yolo9/val',  # Validation 데이터 경로
    'nc': 2,  # 클래스 수
    'names': ['pineapple', 'apple']  # 클래스 이름
}

# yaml 파일 저장 경로 설정
yaml_file_path = '/content/drive/MyDrive/yolo9/data.yaml'

# data.yaml 파일 생성
with open(yaml_file_path, 'w') as yaml_file:
    yaml.dump(data_yaml, yaml_file, default_flow_style=False)

print(f"data.yaml 파일이 {yaml_file_path} 경로에 성공적으로 생성되었습니다.")

 

3. yolo 라이브러리 불러오고 모델 학습 실행

# YOLO 라이브러리 설치 (필요한 경우)
!pip install ultralytics

# YOLO 라이브러리 불러오기
from ultralytics import YOLO

# W&B 비활성화 (필수)
import os
os.environ['WANDB_MODE'] = 'offline'

# YOLOv8 모델 불러오기 (pretrained 모델을 사용할 수 있습니다.)
model = YOLO('yolov8n.pt')  # yolov8n은 가장 가벼운 모델, 필요에 따라 다른 모델로 변경 가능

# 학습 데이터셋 경로 설정 (data.yaml 파일에서 정의한 데이터셋 경로)
data_yaml = '/content/drive/MyDrive/yolo9/data.yaml'

# 모델 학습 실행
model.train(
    data=data_yaml,  # 학습에 사용할 데이터 경로 (data.yaml 파일)
    epochs=100,  # 학습을 반복할 횟수
    imgsz=640,  # 이미지 크기
    batch=16,  # 배치 크기
    name='pineapple_apple_detector',  # 프로젝트 이름
    workers=8  # 데이터 로딩에 사용할 CPU 코어 수
)

-> 머신러닝 실험을 추적하고 관리하는 플랫폼 모델 학습과정을 W&B로 보내지 않겠다는 코드

 

4. 학습된 모델을 명시적으로 저장

# 학습된 모델을 명시적으로 저장하는 코드
# 기본적으로는 YOLOv8이 자동으로 마지막과 가장 성능 좋은 모델을 저장함
model.save('/content/drive/MyDrive/yolo9/trained_model.pth')  # 원하는 경로에 저장

 

5. 학습된 가중치중에 가장 좋은 가중치로 모델을 생성

from ultralytics import YOLO

# Check if 'best.pt' exists in the default runs directory
import os
best_pt_path = '/content/runs/detect/pineapple_apple_detector/weights/best.pt'  # Updated path
if os.path.exists(best_pt_path):
    # Load the model if 'best.pt' exists
    model = YOLO(best_pt_path)
else:
    # If 'best.pt' doesn't exist, try loading the latest weights or the explicitly saved model
    latest_weights_path = '/content/runs/detect/pineapple_apple_detector/weights/last.pt'
    if os.path.exists(latest_weights_path):
        print("Loading the latest weights as 'best.pt' was not found.")
        model = YOLO(latest_weights_path)
    else:
        # Fallback to the explicitly saved model
        print("Loading the explicitly saved model as neither 'best.pt' nor 'last.pt' were found.")
        model = YOLO('/content/drive/MyDrive/yolo9/trained_model.pth')

# Continue with your inference code...

from PIL import Image
import torch

# YOLO 모델을 불러옵니다.
# 모델이 이미 이전 코드에서 로드된 상태여야 합니다.
# 만약 로드되지 않았다면 이전 코드에서처럼 모델을 로드하세요.

# 입력 이미지의 경로를 설정합니다 (경로를 적절히 수정하세요)
image_path = '/content/drive/MyDrive/yolo9/train/11.jpg'

# PIL을 사용하여 이미지를 엽니다
image = Image.open(image_path)

# YOLO 모델을 사용하여 추론을 수행합니다
results = model(image)

# results is a list, access the Results object by indexing
results[0].show()  # 바운딩 박스가 표시된 이미지를 보여줍니다.

from PIL import Image
import torch

# YOLO 모델을 불러옵니다.
# 모델이 이미 이전 코드에서 로드된 상태여야 합니다.
# 만약 로드되지 않았다면 이전 코드에서처럼 모델을 로드하세요.

# 입력 이미지의 경로를 설정합니다 (경로를 적절히 수정하세요)
image_path = '/content/drive/MyDrive/yolo9/train/147.jpg'

# PIL을 사용하여 이미지를 엽니다
image = Image.open(image_path)

# YOLO 모델을 사용하여 추론을 수행합니다
results = model(image)

# results is a list, access the Results object by indexing
results[0].show()  # 바운딩 박스가 표시된 이미지를 보여줍니다.

 

from PIL import Image
import torch

# YOLO 모델을 불러옵니다.
# 모델이 이미 이전 코드에서 로드된 상태여야 합니다.
# 만약 로드되지 않았다면 이전 코드에서처럼 모델을 로드하세요.

# 입력 이미지의 경로를 설정합니다 (경로를 적절히 수정하세요)
image_path = '/content/drive/MyDrive/yolo9/train/11.jpg'

# PIL을 사용하여 이미지를 엽니다
image = Image.open(image_path)

# YOLO 모델을 사용하여 추론을 수행합니다
results = model(image)

# results is a list, access the Results object by indexing
results[0].show()  # 바운딩 박스가 표시된 이미지를 보여줍니다.

 

 

from PIL import Image
import torch

# YOLO 모델을 불러옵니다.
# 모델이 이미 이전 코드에서 로드된 상태여야 합니다.
# 만약 로드되지 않았다면 이전 코드에서처럼 모델을 로드하세요.

# 입력 이미지의 경로를 설정합니다 (경로를 적절히 수정하세요)
image_path = '/content/drive/MyDrive/yolo9/train/147.jpg'

# PIL을 사용하여 이미지를 엽니다
image = Image.open(image_path)

# YOLO 모델을 사용하여 추론을 수행합니다
results = model(image)

# results is a list, access the Results object by indexing
results[0].show()  # 바운딩 박스가 표시된 이미지를 보여줍니다.

import cv2
import torch

# YOLO 모델을 로드합니다 (이미 로드된 상태라면 이 부분은 건너뛰세요)
# 모델을 로드하는 코드
# model = YOLO('best.pt')

# 동영상 파일 경로
video_path = '/content/drive/MyDrive/yolo9/pineapple.mp4'  # 동영상 경로를 적절히 수정하세요
output_video_path = '/content/drive/MyDrive/yolo9/pineapple2.mp4'  # 출력될 동영상 경로

# 동영상을 불러옵니다
cap = cv2.VideoCapture(video_path)

# 동영상의 프레임 크기 및 FPS를 가져옵니다
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)

# 출력 동영상을 저장할 준비를 합니다
fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 코덱 설정
out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break  # 동영상이 끝나면 루프를 종료합니다

    # YOLO 모델을 사용하여 현재 프레임에서 사물 검출을 수행합니다
    results = model(frame)

    # 검출 결과를 현재 프레임에 그립니다
    annotated_frame = results[0].plot()

    # 결과가 그려진 프레임을 출력 동영상에 저장합니다
    out.write(annotated_frame)

    # 결과를 실시간으로 확인하려면 아래 코드 사용 (원하는 경우)
    # cv2.imshow('Detected Frame', annotated_frame)
    # if cv2.waitKey(1) & 0xFF == ord('q'):
    #     break

# 리소스 해제
cap.release()
out.release()
#cv2.destroyAllWindows()

 

trained_model.pt

5.97MB

※ 위 사진과  파일 다운 받고 쥬피터에서 실행

import cv2
import torch

# YOLO 모델을 로드합니다 (이미 로드된 상태라면 이 부분은 건너뛰세요)
# 모델을 로드하는 코드
# model = YOLO('best.pt')

# 동영상 파일 경로
video_path = '/content/drive/MyDrive/yolo9/pineapple.mp4'  # 동영상 경로를 적절히 수정하세요
output_video_path = '/content/drive/MyDrive/yolo9/pineapple2.mp4'  # 출력될 동영상 경로

# 동영상을 불러옵니다
cap = cv2.VideoCapture(video_path)

# 동영상의 프레임 크기 및 FPS를 가져옵니다
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)

# 출력 동영상을 저장할 준비를 합니다
fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 코덱 설정
out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break  # 동영상이 끝나면 루프를 종료합니다

    # YOLO 모델을 사용하여 현재 프레임에서 사물 검출을 수행합니다
    results = model(frame)

    # 검출 결과를 현재 프레임에 그립니다
    annotated_frame = results[0].plot()

    # 결과가 그려진 프레임을 출력 동영상에 저장합니다
    out.write(annotated_frame)

    # 결과를 실시간으로 확인하려면 아래 코드 사용 (원하는 경우)
    # cv2.imshow('Detected Frame', annotated_frame)
    # if cv2.waitKey(1) & 0xFF == ord('q'):
    #     break

# 리소스 해제
cap.release()
out.release()
#cv2.destroyAllWindows()

 

---

■ object detection 을 하기위한 라벨링 작업

 

1. 팬더과 랫서팬더 200장

  수지 : 1.jpg ~  100.jpg
  설현: 101.jpg ~ 200.jpg

2.  라벨링 프로그램으로 라벨링을 합니다. (1시간 정도 소요)

https://cafe.daum.net/oracleoracle/Sl3Y/542

시작 버튼 옆에 검색에 "바이러스 및 위협방지"  검색하세요 ! (백신이 있다면)

 

앱열기 --->  설정 ---> 보호 ---> 실시간 보호 활성화 끄기

 

labelImg-master.zip

10.55MB

 

 

 

 

파일 위치:   C:\Users\itwill\Downloads\labelImg-master\labelImg-master

 

3가지를 누릅니다.

1. opendir 을 누르고 이미지가 있는 디렉토리를 지정

 

2. change save dir  을 누르고 이미지가 있는 디렉토리를 지정

3. yolo 를 선택

pascal VOC 에서 YOLO로 변경

4. 사진을 하나씩 create rectangle 을 누르고  박싱작업하고  suji 로 저장
      설현은 sulhyun 으로 저장

5. 계속 next 누르면서 진행하면됩니다.

6. 다 완료되면 사진마다 좌표정보가 들어있는 text 파일이 있는지 확인
     classes.txt  에  suji  와  sulhyun 만 있는지 확인합니다.

나머지 시간은 시간은 조별 포트폴리오 하면 됩니다. 

*  사물검출 수업  복습 

1.  사물 검출할 대상을 지정
2.  사진을 최소 하나의 사물당 100장씩 
     2개로 분류할거니까 200장 

3. 부지런히 라벨링을 합니다.  (여러시간 소요)

 

 

 

Daum 카페

---

▣ 여러분들이 직접 정성들여서 라벨링한 사진을 학습 시키기 

1. 어제 라벨링한 사진과 정답 텍스트를 준비합니다.

https://cafe.daum.net/oracleoracle/SpOP/342

2. 200장의 사진중에 20장은 val 폴더에 넣고 나머지는 train 폴더에 넣습니다.

  20장중에 10장은 고라니여야하고 또 10장은 너구리여야합니다.

  10장에 정답 텍스트가 같이 있어야합니다.

3. train 과 val 폴더를 각각 압축하고 train.zip 과 val.zip 압축 파일로 생성합니다.
  
4. 코렙에 어제 사용했던 사물 검출 주피터 노트북을 복사를 해서 저장하고
    이름을 변경하세요. 

   yolo8을_이용해서_파인애플과_사과_사물_검출2_잘되는거 최종.ipynb
                                ↓ 
   yolo8을_이용해서_고라니와 너구리_사물_검출2_잘되는거 최종.ipynb

 

 

---

 

1. 구글 드라이브에 val(180장) train(20장) 압축해서 올리기

 

 

2. 코랩에서 압축풀기

from google.colab import drive
drive.mount('/content/drive')

!unzip /content/drive/MyDrive/yolo200/val.zip -d /content/drive/MyDrive/yolo200/

!unzip /content/drive/MyDrive/yolo200/train.zip -d /content/drive/MyDrive/yolo200/

 

3. data.yaml 생성

# 없으면 깔기
!pip install ultralytics

import yaml

# 데이터셋 경로 및 클래스 정보 설정
data_yaml = {
    'train': '/content/drive/MyDrive/yolo200/train',  # Train 데이터 경로
    'val': '/content/drive/MyDrive/yolo200/val',  # Validation 데이터 경로
    'nc': 2,  # 클래스 수
    'names': ['panda', 'lesserpanda']  # 클래스 이름
}

# yaml 파일 저장 경로 설정
yaml_file_path = '/content/drive/MyDrive/yolo200/data.yaml'

# data.yaml 파일 생성
with open(yaml_file_path, 'w') as yaml_file:
    yaml.dump(data_yaml, yaml_file, default_flow_style=False)

print(f"data.yaml 파일이 {yaml_file_path} 경로에 성공적으로 생성되었습니다.")

 

# YOLO 라이브러리 설치 (필요한 경우)
# !pip install ultralytics

# YOLO 라이브러리 불러오기
from ultralytics import YOLO

# W&B 비활성화 (필수)
import os
os.environ['WANDB_MODE'] = 'offline'

# YOLOv8 모델 불러오기 (pretrained 모델을 사용할 수 있습니다.)
model = YOLO('yolov8n.pt')  # yolov8n은 가장 가벼운 모델, 필요에 따라 다른 모델로 변경 가능

# 학습 데이터셋 경로 설정 (data.yaml 파일에서 정의한 데이터셋 경로)
data_yaml = '/content/drive/MyDrive/yolo200/data.yaml'

# 모델 학습 실행
model.train(
    data=data_yaml,  # 학습에 사용할 데이터 경로 (data.yaml 파일)
    epochs=100,  # 학습을 반복할 횟수
    imgsz=640,  # 이미지 크기
    batch=16,  # 배치 크기
    name='panda_detector',  # 프로젝트 이름
    workers=8  # 데이터 로딩에 사용할 CPU 코어 수
)

# 학습된 모델을 명시적으로 저장하는 코드
# 기본적으로는 YOLOv8이 자동으로 마지막과 가장 성능 좋은 모델을 저장함
model.save('/content/drive/MyDrive/yolo200/trained_model.pt')  # 원하는 경로에 저장

trained_model.pt

5.96MB

 

 

 

 

import cv2
import numpy as np
from ultralytics import YOLO  # YOLOv8 라이브러리 임포트
from tkinter import filedialog
from tkinter import *
from PIL import ImageTk, Image
import matplotlib.pyplot as plt

LABELS = ['panda', 'lesserpanda']

# YOLOv8 모델 로드
model = YOLO('C:\\imageall2\\trained_model.pt')  # YOLOv8 학습된 가중치 파일 경로

CONFIDENCE = 0.3 #팬더와 랫서판더일 확률 30% 이상일 때만 사물검출하겠다
THRESHOLD = 0.3  # NMS(비최대 억제 임계값) 여러개의 박스들에서 사물이 있을 만한 확률이 30% 이상인 것들만 사물 검출

def main(img_path):
    img = cv2.imread(img_path)
    
    # YOLOv8을 사용하여 이미지에서 객체 검출
    results = model.predict(img, conf=CONFIDENCE, iou=THRESHOLD)

    # YOLOv8의 결과로 바운딩 박스 그리기
    annotated_img = results[0].plot()

    # YOLOv8 결과 이미지를 저장
    cv2.imwrite('C:\\imageall2\\detection12.png', annotated_img)

##### tkinter GUI 부분 #####

# tkinter GUI 창 생성
top = Tk()
top.geometry('700x600')
top.title('Animal Object Detection')
top.configure(background='snow2')

# 기본 이미지 로드
img = Image.open('c:\\imageall2\\1.jpg')
img = img.resize((400, 320))
im = ImageTk.PhotoImage(img, master=top)
sign_image = Label(top, image=im)
sign_image.place(relx=0.28, rely=0.1)

##### 화면 타이틀 #####
heading = Label(top, text="Panda Classifier", pady=20, font=('CG Omega', 20, 'bold'))
heading.configure(background='snow2', foreground='navy')
heading.pack()

##### 이미지를 창에 업로드하는 함수 #####
def upload_image():
    plt.clf()  # 현재 그림을 지움
    try:
        global file_path
        file_path = filedialog.askopenfilename()  # 파일 선택창 열기
        
        uploaded = Image.open(file_path)
        uploaded = uploaded.resize((400, 320))
        im = ImageTk.PhotoImage(uploaded)
        sign_image.configure(image=im)
        sign_image.image = im
    except:
        pass

##### 사진 속에 사물 검출하는 함수 #####
def detect_image():
    plt.clf()  # 현재 그림을 지움
    try:
        main(file_path)  # YOLOv8 객체 검출 실행
        
        # 검출된 이미지를 로드하여 tkinter 창에 표시
        file_path2 = 'C:\\imageall2\\detection12.png'
        uploaded2 = Image.open(file_path2)
        uploaded2 = uploaded2.resize((400, 320))
        
        im2 = ImageTk.PhotoImage(uploaded2)
        sign_image.configure(image=im2)
        sign_image.image = im2
    except:
        pass

##### 업로드 버튼 생성 #####
upload = Button(top, text="Upload image", relief=RIDGE, command=upload_image, padx=10, pady=5)
upload.place(relx=0.2, rely=0.9)
upload.configure(background='mediumpurple4', foreground='white', font=('arial', 12, 'bold'))

##### 객체 검출 버튼 생성 #####
upload2 = Button(top, text="Detect image", relief=RIDGE, command=detect_image, padx=10, pady=5)
upload2.place(relx=0.6, rely=0.9)
upload2.configure(background='mediumpurple4', foreground='white', font=('arial', 12, 'bold'))

sign_image.pack(side=BOTTOM, expand=True)

# tkinter 창 실행
top.mainloop()

 

 

 

※ 동영상 인식(코렙)

import cv2
import torch
from ultralytics import YOLO

# YOLO 모델을 로드합니다 (이미 로드된 상태라면 이 부분은 건너뛰세요)
# 모델을 로드하는 코드
model = YOLO('/content/drive/MyDrive/yolo200/trained_model.pt')

# 동영상 파일 경로
video_path = '/content/drive/MyDrive/yolo200/pandamv.mkv'  # 동영상 경로를 적절히 수정하세요
output_video_path = '/content/drive/MyDrive/yolo200/pandamv_output.mp4'  # 출력될 동영상 경로

# 동영상을 불러옵니다
cap = cv2.VideoCapture(video_path)

# 동영상의 프레임 크기 및 FPS를 가져옵니다
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)

# 출력 동영상을 저장할 준비를 합니다
fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 코덱 설정
out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break  # 동영상이 끝나면 루프를 종료합니다

    # YOLO 모델을 사용하여 현재 프레임에서 사물 검출을 수행합니다
    results = model(frame)

    # 검출 결과를 현재 프레임에 그립니다
    annotated_frame = results[0].plot()

    # 결과가 그려진 프레임을 출력 동영상에 저장합니다
    out.write(annotated_frame)

    # 결과를 실시간으로 확인하려면 아래 코드 사용 (원하는 경우)
    # cv2.imshow('Detected Frame', annotated_frame)
    # if cv2.waitKey(1) & 0xFF == ord('q'):
    #     break

# 리소스 해제
cap.release()
out.release()
#cv2.destroyAllWindows()

 

'''

 

정확도를 올리는 총코드

import cv2
import torch
from ultralytics import YOLO

# YOLO 모델을 로드합니다 (이미 로드된 상태라면 이 부분은 건너뛰세요)
# 모델을 로드하는 코드
model = YOLO('/content/drive/MyDrive/yolo200/trained_model.pt')

# 동영상 파일 경로
video_path = '/content/drive/MyDrive/yolo200/pandamv.mkv'  # 동영상 경로를 적절히 수정하세요
output_video_path = '/content/drive/MyDrive/yolo200/pandamv_output2.mp4'  # 출력될 동영상 경로

# 동영상을 불러옵니다
cap = cv2.VideoCapture(video_path)

# 동영상의 프레임 크기 및 FPS를 가져옵니다
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)

# 출력 동영상을 저장할 준비를 합니다
fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 코덱 설정
out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))

# Confidence와 IOU 임계값을 설정
confidence_threshold = 0.6  # 원하는 신뢰도 임계값으로 조정
iou_threshold = 0.5         # 원하는 IOU 임계값으로 조정

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break  # 동영상이 끝나면 루프를 종료합니다

    # YOLO 모델을 사용하여 현재 프레임에서 사물 검출을 수행합니다
    results = model(frame)

    # 검출 결과를 현재 프레임에 그립니다
    annotated_frame = results[0].plot()

    # 결과가 그려진 프레임을 출력 동영상에 저장합니다
    out.write(annotated_frame)

    # 결과를 실시간으로 확인하려면 아래 코드 사용 (원하는 경우)
    # cv2.imshow('Detected Frame', annotated_frame)
    # if cv2.waitKey(1) & 0xFF == ord('q'):
    #     break

# 리소스 해제
cap.release()
out.release()
#cv2.destroyAllWindows()