fortnite-game-coach / src /vidprocessing.py
reeeeemo
Added YOLO model
019f9fc
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))
# read frame_count / interval of frames
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}')
# write photo to file
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)
### PRETRAINED MODEL DETECTION CODE
yolo = ObjectDetector(pretrained_model=(base_dir / 'best.pt'), debug=True)
yolo.detect_object(frames)
### TRAINING CODE FOR YOLO MODEL
# yolo = ObjectDetector()
# yolo.train_model(base_dir / 'data' / 'model_data' / 'fortnite_train.yaml')
### VIDEO FRAME CUTTING CODE
# create_images_of_video(vid)
### MODEL EXPORT + VIDEO RELEASE
yolo.export_model()
vid.release()