AlexNet_CNN_Visualisation / video_preprocessing.py
makiisthebes's picture
Upload 18 files
4ec6f12 verified
# Taking each video cropping to 520, 400, 3, and obtaining 12 key frames for the usage in a dataset.
# Frames will be saved as individual images.
import numpy as np
import cv2, os
from tqdm import tqdm
# Obtain 12 different keyframes from the video, equally spaced.
# Actually to increase dataset we will just increase the number of frames allowed to say 50 per video. 900 images for dataset, 200 are for training.
def get_equal_elements(array, num_elements=12):
"""
Takes a specific number of elements equally spaced from an array.
Args:
array: The input array.
num_elements: The number of elements to take (default 12).
Returns:
A list of elements from the array.
"""
if num_elements > len(array):
print(f"Number of elements cannot be greater than array length : {len(array)}")
return []
step_size = len(array) // (num_elements - 1) # Avoid extra element with floor division
return array[::step_size] # Slice with step size
def video_to_keyframes(video_filename):
cap = cv2.VideoCapture(video_filename)
frames = []
while (cap.isOpened()):
ret, frame = cap.read()
try:
# 720, 1280, 3
# Remove top 200 pixels from video frame.
# print((frame.shape)) # (720, 1280, 3)
frame = frame[200:, 440:840]
# print((frame.shape)) # (520, 400, 3)
frames.append(frame)
# cv2.imshow('frame',frame)
except Exception as e:
print(f"Error is {e}")
if frame == None:
break
continue
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
print("Done obtaining captured frames.")
selected_frames = get_equal_elements(frames, num_elements=50)
print("Obtained selected frames.")
# Extract filename without extension for directory creation
filename_no_ext = video_filename.split('.')[0]
if "left" in filename_no_ext:
filename_no_ext = "left"
else:
filename_no_ext = "right"
# Create directory for the video if it doesn't exist
try:
os.makedirs(os.path.join("dataset", filename_no_ext))
except FileExistsError:
pass # Directory already exists, ignore
# Save each selected frame as an image with filename_frame_number.jpg format
# Create images under directory for that specific source video file name, like left for left1.mp4.
for i, frame in enumerate(selected_frames):
print(filename_no_ext)
file_name = f"{video_filename.split('.')[0].split('/')[-1]}_frame_{i}.jpg"
print(file_name)
image_path = os.path.join("dataset", filename_no_ext, file_name)
print(f"Write to disk. {image_path}")
print("Resized to 224,224")
target_height, target_width = 224, 224
if type(frame) != type(None):
frame = resize_with_aspect_ratio(frame, target_height, target_width)
# print(type(frame), frame)
cv2.imwrite(image_path, frame)
else:
continue
print("Saved images for each all selected frames.")
return selected_frames
# Resize images.
# Target dimensions for AlexNet
# Resize the image with aspect ratio preservation
def resize_with_aspect_ratio(image, target_height, target_width):
height, width = image.shape[:2]
if height == target_height and width == target_width:
return image
if height > width:
new_width = int(width * (target_height / height))
# resized_image = cv2.resize(image, (new_width, target_height))
resized_image = cv2.resize(image, (target_width, target_height))
return resized_image
else:
new_height = int(height * (target_width / width))
resized_image = cv2.resize(image, (target_width, new_height))
# Crop the center of the resized image to match target dimensions
start_x = int((resized_image.shape[1] - target_width) / 2)
start_y = int((resized_image.shape[0] - target_height) / 2)
return resized_image[start_y:start_y + target_height, start_x:start_x + target_width]
if __name__ == "__main__":
# Load the image and resize
BASE_PATH = "dataset/src/"
videos = os.listdir(BASE_PATH)
# print(videos)
target_height, target_width = 224, 224
for video_file in tqdm(videos):
selected_frames = video_to_keyframes(os.path.join(BASE_PATH, video_file))
# image = cv2.imread(image_path) # Or use your image data
# resized_image = resize_with_aspect_ratio(image, target_height, target_width)