diff --git a/DeepEncode.py b/DeepEncode.py index 6ae73d9..7a033e2 100644 --- a/DeepEncode.py +++ b/DeepEncode.py @@ -3,6 +3,7 @@ import os from featureExtraction import preprocess_frame +from globalVars import PRESET_SPEED_CATEGORIES os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1' @@ -14,6 +15,8 @@ from video_compression_model import VideoCompressionModel # Constants COMPRESSED_VIDEO_FILE = 'compressed_video.avi' MAX_FRAMES = 0 # Limit the number of frames processed +CRF = 51 +SPEED = PRESET_SPEED_CATEGORIES.index("ultrafast") # Load the trained model MODEL = tf.keras.models.load_model('models/model.tf', custom_objects={'VideoCompressionModel': VideoCompressionModel}) @@ -37,7 +40,7 @@ def predict_frame(uncompressed_frame): #display_frame = np.clip(cv2.cvtColor(uncompressed_frame, cv2.COLOR_BGR2RGB) * 255.0, 0, 255).astype(np.uint8) #cv2.imshow("uncomp", uncompressed_frame) - frame = preprocess_frame(uncompressed_frame) + frame = preprocess_frame(uncompressed_frame, CRF, SPEED) compressed_frame = MODEL.predict([np.expand_dims(frame, axis=0)])[0] @@ -54,7 +57,7 @@ height, width = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(cap.get(cv2.CAP_PRO cap.release() fourcc = cv2.VideoWriter_fourcc(*'XVID') -out = cv2.VideoWriter(COMPRESSED_VIDEO_FILE, fourcc, 24.0, (width, height)) +out = cv2.VideoWriter(COMPRESSED_VIDEO_FILE, fourcc, 24.0, (width, height), True) if not out.isOpened(): print("Error: VideoWriter could not be opened.") diff --git a/featureExtraction.py b/featureExtraction.py index ce8168c..664ff0d 100644 --- a/featureExtraction.py +++ b/featureExtraction.py @@ -3,7 +3,7 @@ import cv2 import numpy as np -from globalVars import HEIGHT, WIDTH +from globalVars import HEIGHT, NUM_PRESET_SPEEDS, WIDTH def extract_edge_features(frame): @@ -39,17 +39,23 @@ def extract_histogram_features(frame, bins=64): return np.array(feature_vector) -def preprocess_frame(frame): +def preprocess_frame(frame, crf, speed): # Check frame dimensions and resize if necessary if frame.shape[:2] != (HEIGHT, WIDTH): frame = cv2.resize(frame, (WIDTH, HEIGHT), interpolation=cv2.INTER_NEAREST) - # Extract features - edge_feature = extract_edge_features(frame) - histogram_feature = extract_histogram_features(frame) - - histogram_feature_image = np.full((HEIGHT, WIDTH), histogram_feature.mean()) # Convert histogram feature to image-like shape - combined_feature = np.stack([edge_feature, histogram_feature_image], axis=-1) + # Scale frame to [0, 1] + compressed_frame = frame / 255.0 - compressed_frame = frame / 255.0 # Assuming the frame is uint8, scale to [0, 1] - return compressed_frame \ No newline at end of file + # Scale CRF and SPEED to [0, 1] (assuming they are within known bounds) + crf_scaled = crf / 51 + speed_scaled = speed / NUM_PRESET_SPEEDS + + # Create images with the CRF and SPEED values, filling extra channels + crf_image = np.full((HEIGHT, WIDTH, 1), crf_scaled) # Note the added dimension + speed_image = np.full((HEIGHT, WIDTH, 1), speed_scaled) # Note the added dimension + + # Combine the frames with the CRF and SPEED images + combined_frame = np.concatenate([compressed_frame, crf_image, speed_image], axis=-1) + + return combined_frame diff --git a/video_compression_model.py b/video_compression_model.py index c68609a..69f9c0a 100644 --- a/video_compression_model.py +++ b/video_compression_model.py @@ -5,7 +5,7 @@ import cv2 import numpy as np import tensorflow as tf from featureExtraction import preprocess_frame -from globalVars import HEIGHT, LOGGER, NUM_COLOUR_CHANNELS, WIDTH +from globalVars import HEIGHT, LOGGER, NUM_COLOUR_CHANNELS, NUM_PRESET_SPEEDS, PRESET_SPEED_CATEGORIES, WIDTH #from tensorflow.keras.mixed_precision import Policy @@ -22,13 +22,16 @@ def data_generator(videos, batch_size): video_path = os.path.join(os.path.dirname("test_data/validation/validation.json"), video_details["compressed_video_file"]) uncompressed_video_path = os.path.join(os.path.dirname("test_data/validation/validation.json"), video_details["original_video_file"]) + CRF = video_details["crf"] / 51 + SPEED = PRESET_SPEED_CATEGORIES.index(video_details["preset_speed"]) + # Open the video files cap_compressed = cv2.VideoCapture(video_path) cap_uncompressed = cv2.VideoCapture(uncompressed_video_path) # Lists to store the processed frames - compressed_frame_batch = [] # Input data (Training) - uncompressed_frame_batch = [] # Target data (Target) + compressed_frame_batch = [] # Input data (Target) + uncompressed_frame_batch = [] # Target data (Training) # Read and process frames from both videos while cap_compressed.isOpened() and cap_uncompressed.isOpened(): @@ -37,11 +40,11 @@ def data_generator(videos, batch_size): if not ret_compressed or not ret_uncompressed: break - # Preprocess the compressed frame (input) - compressed_frame = preprocess_frame(compressed_frame) + # Preprocess the compressed frame (target) + compressed_frame = preprocess_frame(compressed_frame, CRF, SPEED) - # Preprocess the uncompressed frame (target) - uncompressed_frame = preprocess_frame(uncompressed_frame) # Modify if different preprocessing is needed for target frames + # Preprocess the uncompressed frame (input) + uncompressed_frame = preprocess_frame(uncompressed_frame, 0, PRESET_SPEED_CATEGORIES.index("veryslow")) # Modify if different preprocessing is needed for target frames # Append processed frames to batches compressed_frame_batch.append(compressed_frame) @@ -49,7 +52,7 @@ def data_generator(videos, batch_size): # If batch is complete, yield it if len(compressed_frame_batch) == batch_size: - yield (np.array(compressed_frame_batch), np.array(uncompressed_frame_batch)) # Yielding Training and Target data + yield (np.array(uncompressed_frame_batch), np.array(compressed_frame_batch)) # Yielding Training and Target data compressed_frame_batch = [] uncompressed_frame_batch = [] @@ -59,7 +62,7 @@ def data_generator(videos, batch_size): # If there are frames left that don't fill a whole batch, send them anyway if len(compressed_frame_batch) > 0: - yield (np.array(compressed_frame_batch), np.array(uncompressed_frame_batch)) + yield (np.array(uncompressed_frame_batch), np.array(compressed_frame_batch)) class VideoCompressionModel(tf.keras.Model): def __init__(self): @@ -67,7 +70,7 @@ class VideoCompressionModel(tf.keras.Model): LOGGER.debug("Initializing VideoCompressionModel.") # Input shape (includes channels for edges and histogram) - input_shape_with_histogram = (HEIGHT, WIDTH, NUM_COLOUR_CHANNELS) + input_shape_with_histogram = (HEIGHT, WIDTH, NUM_COLOUR_CHANNELS + 2) # Encoder part of the model self.encoder = tf.keras.Sequential([