import torch
import os
import cv2
from tqdm import tqdm
model = "yolov5"
# yolov5l model 사용
model.conf = 0.25
model.classes = [1]
path = "test/images/path"
file_list = os.listdir(path)
# path 에 있는 파일들을 전부 불러옴
for idx in tqdm(iterable=range(len(file_list))):
# list 안에 있는 파일들을 하나씩 확인
if file_list[idx].split(".")[-1] == "mp4":
# 파일이 mp4 확장자 (영상) 인 경우
vidcap = cv2.VideoCapture(path + file_list[idx])
# cv2.VideoCapture 인스턴스 생성
fourcc = cv2.VideoWriter_fourcc(*'DIVX')
# video 를 저장하기 위해 필요한 코덱 포맷
w = round(vidcap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = round(vidcap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = vidcap.get(cv2.CAP_PROP_FPS)
# 원본 영상의 width, height, fps 를 가져옴
out = cv2.VideoWriter("video/save/path", fourcc, fps, (w,h))
# cv2.VideoWriter 인스턴스 생성
while True:
# frame capture 시작
try:
_, frame = vidcap.read()
results = model(frame)
# 캡처한 프레임을 yolov5l 에 넣어 inference 함
coords = results.xyxy[0]
# inference 결과물로 bbox 의 좌표, class, name 출력
if len(coords) == 0:
# bbox 가 없을 때 bbox 가 없는 프레임 저장
out.write(frame)
else:
for i in range(len(coords)):
try:
x1, y1, x2, y2 = int(coords[i][0]), int(coords[i][1]), int(coords[i][2]), int(coords[i][3])
w = x2 - x1
h = y2 - y1
center_x = int(x1 + w / 2)
center_y = int(y1 + h / 2)
img = cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
# results 값의 좌표값들로 프레임에 사각형을 그림
img = cv2.line(img, (center_x, center_y), (center_x, center_y), (0,0,0), 4)
# results 값의 좌표값들로 검출 된 객체의 중심좌표에 점을 찍음
out.write(img)
# 영상 저장
except:
pass
except:
break
vidcap.release()
elif file_list[idx].split(".")[-1] == "jpg" or "png" or "jpeg":
img = cv2.imread(path + file_list[idx])
results = model(img)
coords = results.xyxy[0]
...
