Black and white

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.")

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
+    # Scale frame to [0, 1]
-    edge_feature = extract_edge_features(frame)
+    compressed_frame = frame / 255.0
-    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)
 
-    compressed_frame = frame / 255.0  # Assuming the frame is uint8, scale to [0, 1]
+    # Scale CRF and SPEED to [0, 1] (assuming they are within known bounds)
-    return compressed_frame
+    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

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)
+            compressed_frame_batch = [] # Input data (Target)
-            uncompressed_frame_batch = [] # Target 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)
+                # Preprocess the compressed frame (target)
-                compressed_frame = preprocess_frame(compressed_frame)
+                compressed_frame = preprocess_frame(compressed_frame, CRF, SPEED)
                 
-                # Preprocess the uncompressed frame (target)
+                # Preprocess the uncompressed frame (input)
-                uncompressed_frame = preprocess_frame(uncompressed_frame) # Modify if different preprocessing is needed for target frames
+                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([