Black and white
This commit is contained in:
parent
93ef52e66f
commit
e7af02cb4f
3 changed files with 34 additions and 22 deletions
|
@ -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.")
|
||||
|
|
|
@ -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
|
||||
# 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
|
||||
|
|
|
@ -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([
|
||||
|
|
Reference in a new issue