|
import cv2 |
|
from typing import List |
|
from os.path import dirname, abspath |
|
from pathlib import Path |
|
import numpy as np |
|
from detectobjects import ObjectDetector |
|
import yaml |
|
|
|
base_dir = Path(dirname(dirname(abspath(__file__)))) |
|
|
|
def get_video(path: str) -> cv2.VideoCapture: |
|
video = cv2.VideoCapture(path) |
|
if not video.isOpened(): |
|
raise ValueError(f'Could not open video file: {path}') |
|
|
|
fps = video.get(cv2.CAP_PROP_FPS) |
|
frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) |
|
duration = frame_count/fps |
|
|
|
print(f'FPS: {fps}\nFrame Count: {frame_count}\nDuration: {duration}') |
|
|
|
return video |
|
|
|
def get_frames(video: cv2.VideoCapture, frame_start: int, frame_end: int) -> List[np.ndarray]: |
|
|
|
frames = [] |
|
for i in range(frame_start, frame_end+1): |
|
video.set(cv2.CAP_PROP_POS_FRAMES, i) |
|
ret, frame = video.read() |
|
if not ret: |
|
raise ValueError(f'Could not read frame {i}') |
|
|
|
frames.append(frame) |
|
|
|
return frames |
|
|
|
|
|
def create_images_of_video(video: cv2.VideoCapture, interval: int = 100): |
|
frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) |
|
|
|
|
|
for i in range(0, frame_count, interval): |
|
video.set(cv2.CAP_PROP_POS_FRAMES, i) |
|
success, frame = video.read() |
|
if not success: |
|
raise ValueError(f'Could not read frame {i}') |
|
|
|
|
|
cv2.imwrite(base_dir / "data" / 'model_data' / 'temp_vid_folder' /f'{i}.png', frame) |
|
|
|
|
|
if __name__ == "__main__": |
|
vid = get_video(base_dir / "data" / "video_data" / "fortnite_remo_three.mp4") |
|
frames = get_frames(vid, 21100, 21110) |
|
|
|
|
|
yolo = ObjectDetector(pretrained_model=(base_dir / 'best.pt'), debug=True) |
|
yolo.detect_object(frames) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
yolo.export_model() |
|
vid.release() |