data_augmentation.py

import cv2
import numpy as np

def flip(input_path, output_path):
    """Open video specified by 'input_path', save flipped video to 'output_path'"""
    cap = cv2.VideoCapture(input_path)
    if not cap.isOpened():
        print("Fail to load video")
        return
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        flipped_frame = cv2.flip(frame, 1)
        out.write(flipped_frame)
    cap.release()
    out.release()

def modify_brightness(input_path, output_path, value):
    """Open video specified by 'input_path', modify saturation by 'value', and save to 'output_path'\n
    value > 0 => brighter, else darker"""
    cap = cv2.VideoCapture(input_path)
    if not cap.isOpened():
        print("Fail to load video")
        return
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.add(frame, value) 
        out.write(frame)
    cap.release()
    out.release()

def rotate(input_path, output_path, angle):
    """Open video specified by 'input_path', rotate an 'angle', and save to 'output_path'\n
    angle > 0 => counter clockwise, else clockwise"""
    cap = cv2.VideoCapture(input_path)
    if not cap.isOpened():
        print("Fail to load video")
        return
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.warpAffine(frame, cv2.getRotationMatrix2D((width / 2, height / 2), angle, 1), (width, height))
        out.write(frame)
    cap.release()
    out.release()

def add_gaussian_noise(input_path, output_path, intensity):
    """Open video specified by 'input_path', add noise with 'intensity', and save to 'output_path'\n
    intensity varies between [0,1), 0 means noise free"""
    cap = cv2.VideoCapture(input_path)
    if not cap.isOpened():
        print("Fail to load video")
        return
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        frame = frame/255
        frame = cv2.add(frame,np.random.randn(width,height,3)*(intensity**0.5)*0.1)
        frame = frame*255
        out.write(frame.astype(np.uint8))
    cap.release()
    out.release()