src/vidprocessing.py

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()