jordon/DeepEncode
Archived
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Black and white

Browse source
This commit is contained in:
Jordon Brooks 2023-08-13 18:53:21 +01:00
parent 93ef52e66f
commit e7af02cb4f
3 changed files with 34 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

23
video_compression_model.py
View file

@ -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([