68 lines
2.4 KiB
Python
68 lines
2.4 KiB
Python
import tensorflow as tf
|
|
import numpy as np
|
|
import cv2
|
|
from video_compression_model import VideoCompressionModel
|
|
|
|
# Constants
|
|
CHUNK_SIZE = 24 # Adjust based on available memory and video resolution
|
|
COMPRESSED_VIDEO_FILE = 'compressed_video.mp4'
|
|
MAX_FRAMES = 24 # Limit the number of frames processed
|
|
|
|
# Load the trained model
|
|
model = tf.keras.models.load_model('models/model.keras', custom_objects={'VideoCompressionModel': VideoCompressionModel})
|
|
|
|
# Load the uncompressed video
|
|
UNCOMPRESSED_VIDEO_FILE = 'test_data/training_video.mkv'
|
|
|
|
def load_frame_from_video(video_file, frame_num):
|
|
cap = cv2.VideoCapture(video_file)
|
|
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_num)
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
return None
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 # Normalize and convert to float32
|
|
cap.release()
|
|
|
|
#display_frame = np.clip(frame * 255.0, 0, 255).astype(np.uint8)
|
|
#cv2.imshow("uncomp", display_frame)
|
|
#cv2.waitKey(0) # Add this line to hold the display window until a key is pressed
|
|
|
|
|
|
return frame
|
|
|
|
def predict_frame(uncompressed_frame, model, crf_value, preset_speed_value):
|
|
crf_array = np.array([crf_value])
|
|
preset_speed_array = np.array([preset_speed_value])
|
|
|
|
compressed_frame = model.predict({
|
|
"frame": np.array([uncompressed_frame]),
|
|
"crf": crf_array,
|
|
"preset_speed": preset_speed_array
|
|
})
|
|
return compressed_frame[0]
|
|
|
|
cap = cv2.VideoCapture(UNCOMPRESSED_VIDEO_FILE)
|
|
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
cap.release()
|
|
|
|
if MAX_FRAMES != 0 and total_frames > MAX_FRAMES:
|
|
total_frames = MAX_FRAMES
|
|
|
|
crf_value = 25.0 # Example CRF value
|
|
preset_speed_value = 2 # Index for "fast" in our defined list
|
|
|
|
height, width = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
fourcc = cv2.VideoWriter_fourcc(*'H264')
|
|
out = cv2.VideoWriter(COMPRESSED_VIDEO_FILE, fourcc, 24.0, (width, height))
|
|
|
|
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_value, preset_speed_value)
|
|
|
|
compressed_frame = np.clip(compressed_frame * 255.0, 0, 255).astype(np.uint8)
|
|
compressed_frame = cv2.cvtColor(compressed_frame, cv2.COLOR_RGB2BGR)
|
|
out.write(compressed_frame)
|
|
cv2.imshow("output", compressed_frame)
|
|
|
|
out.release()
|
|
print("Compression completed.")
|