90 lines
3.1 KiB
Python
90 lines
3.1 KiB
Python
# DeepEncode.py
|
|
|
|
import os
|
|
|
|
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
|
|
|
|
import tensorflow as tf
|
|
import numpy as np
|
|
import cv2
|
|
from video_compression_model import PRESET_SPEED_CATEGORIES, VideoCompressionModel
|
|
|
|
# Constants
|
|
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 = 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})
|
|
|
|
# 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()
|
|
|
|
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])
|
|
|
|
crf_array = np.expand_dims(np.array([crf_value]), axis=-1) # Shape: (1, 1)
|
|
preset_speed_array = np.expand_dims(np.array([preset_speed_value]), axis=-1) # Shape: (1, 1)
|
|
|
|
|
|
# Expand dimensions to include batch size
|
|
uncompressed_frame = np.expand_dims(uncompressed_frame, 0)
|
|
|
|
#display_frame = np.clip(cv2.cvtColor(uncompressed_frame[0], cv2.COLOR_BGR2RGB) * 255.0, 0, 255).astype(np.uint8)
|
|
#cv2.imshow("uncomp", display_frame)
|
|
#cv2.waitKey(0)
|
|
|
|
compressed_frame = model.predict({
|
|
"compressed_frame": uncompressed_frame,
|
|
"uncompressed_frame": uncompressed_frame,
|
|
"crf": crf_array,
|
|
"preset_speed": preset_speed_array
|
|
})
|
|
|
|
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(1)
|
|
|
|
return compressed_frame[0]
|
|
|
|
cap = cv2.VideoCapture(UNCOMPRESSED_VIDEO_FILE)
|
|
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
height, width = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
cap.release()
|
|
|
|
fourcc = cv2.VideoWriter_fourcc(*'XVID')
|
|
out = cv2.VideoWriter(COMPRESSED_VIDEO_FILE, fourcc, 24.0, (width, height))
|
|
|
|
if not out.isOpened():
|
|
print("Error: VideoWriter could not be opened.")
|
|
exit()
|
|
|
|
if MAX_FRAMES != 0 and total_frames > MAX_FRAMES:
|
|
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_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)
|
|
out.write(compressed_frame)
|
|
#cv2.imshow("output", compressed_frame)
|
|
|
|
out.release()
|
|
print("Compression completed.")
|