This commit is contained in:
Jordon Brooks 2023-07-30 17:29:32 +01:00
parent 60c6c97071
commit 93ccce5ec1
6 changed files with 46 additions and 34 deletions

View file

@ -3,14 +3,14 @@
import tensorflow as tf
import numpy as np
import cv2
from video_compression_model import VideoCompressionModel
from video_compression_model import PRESET_SPEED_CATEGORIES, VideoCompressionModel
# Constants
CHUNK_SIZE = 10 # Adjust based on available memory and video resolution
CHUNK_SIZE = 24 # Adjust based on available memory and video resolution
COMPRESSED_VIDEO_FILE = 'compressed_video.avi'
MAX_FRAMES = 0 # Limit the number of frames processed
CRF = 25.0 # Example CRF value
PRESET_SPEED = 4 # Index for "fast" in our defined list
CRF = 24.0 # Example CRF value
PRESET_SPEED = "veryslow" # Index for "fast" in our defined list
# Load the trained model
model = tf.keras.models.load_model('models/model.tf', custom_objects={'VideoCompressionModel': VideoCompressionModel})
@ -42,6 +42,7 @@ def predict_frame(uncompressed_frame, model, crf_value, preset_speed_value):
compressed_frame = model.predict({
"compressed_frame": uncompressed_frame,
"uncompressed_frame": uncompressed_frame,
"crf": crf_array,
"preset_speed": preset_speed_array
})
@ -49,7 +50,7 @@ def predict_frame(uncompressed_frame, model, crf_value, preset_speed_value):
display_frame = np.clip(cv2.cvtColor(compressed_frame[0], cv2.COLOR_BGR2RGB) * 255.0, 0, 255).astype(np.uint8)
cv2.imshow("comp", display_frame)
cv2.waitKey(10)
cv2.waitKey(1)
return compressed_frame[0]
@ -70,7 +71,7 @@ if MAX_FRAMES != 0 and total_frames > MAX_FRAMES:
for i in range(total_frames):
uncompressed_frame = load_frame_from_video(UNCOMPRESSED_VIDEO_FILE, frame_num=i)
compressed_frame = predict_frame(uncompressed_frame, model, CRF, PRESET_SPEED)
compressed_frame = predict_frame(uncompressed_frame, model, CRF, PRESET_SPEED_CATEGORIES.index(PRESET_SPEED))
compressed_frame = np.clip(compressed_frame * 255.0, 0, 255).astype(np.uint8)
compressed_frame = cv2.cvtColor(compressed_frame, cv2.COLOR_RGB2BGR)