64 lines
2 KiB
Python
64 lines
2 KiB
Python
import tensorflow as tf
|
|
import numpy as np
|
|
import cv2
|
|
from video_compression_model import VideoCompressionModel
|
|
|
|
# Constants
|
|
NUM_CHANNELS = 3
|
|
|
|
# Step 2: Load the trained model
|
|
model = tf.keras.models.load_model('models/model.keras', custom_objects={'VideoCompressionModel': VideoCompressionModel})
|
|
|
|
# Step 3: Load the uncompressed video
|
|
UNCOMPRESSED_VIDEO_FILE = 'test_data/test_video.mkv'
|
|
|
|
def load_frames_from_video(video_file, num_frames = 0):
|
|
print("Extracting video frames...")
|
|
cap = cv2.VideoCapture(video_file)
|
|
frames = []
|
|
count = 0
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
print("Max frames from file reached")
|
|
break
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
frames.append(frame)
|
|
count += 1
|
|
if num_frames == 0 or count >= num_frames:
|
|
print("Max Frames wanted reached: ", num_frames)
|
|
break
|
|
cap.release()
|
|
print("Extraction Complete")
|
|
return frames
|
|
|
|
uncompressed_frames = load_frames_from_video(UNCOMPRESSED_VIDEO_FILE, 200)
|
|
if len(uncompressed_frames) == 0 or None:
|
|
print("IO ERROR!")
|
|
exit()
|
|
|
|
uncompressed_frames = np.array(uncompressed_frames) / 255.0
|
|
|
|
if len(uncompressed_frames) == 0 or None:
|
|
print("np.array ERROR!")
|
|
exit()
|
|
|
|
# Step 4: Compress the video frames using the loaded model
|
|
compressed_frames = model.predict(uncompressed_frames)
|
|
|
|
# Step 5: Save the compressed video frames
|
|
COMPRESSED_VIDEO_FILE = 'compressed_video.mkv'
|
|
|
|
def save_frames_as_video(frames, video_file):
|
|
print("Saving video frames...")
|
|
height, width = frames[0].shape[:2]
|
|
fourcc = cv2.VideoWriter_fourcc(*'XVID')
|
|
out = cv2.VideoWriter(video_file, fourcc, 24.0, (width, height))
|
|
for frame in frames:
|
|
frame = np.clip(frame * 255.0, 0, 255).astype(np.uint8)
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
|
|
out.write(frame)
|
|
out.release()
|
|
|
|
save_frames_as_video(compressed_frames, COMPRESSED_VIDEO_FILE)
|
|
print("Compression completed.")
|