Upload data_augmentation.py

data_augmentation.py

@@ -0,0 +1,83 @@
+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()