88 lines
2.7 KiB
Python
88 lines
2.7 KiB
Python
# DeepEncode.py
|
|
|
|
import os
|
|
|
|
from featureExtraction import preprocess_frame, psnr, scale_crf, scale_speed_preset
|
|
from globalVars import PRESET_SPEED_CATEGORIES
|
|
|
|
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
|
|
|
|
import tensorflow as tf
|
|
import numpy as np
|
|
import cv2
|
|
from video_compression_model import VideoCompressionModel, combine_batch
|
|
|
|
# Constants
|
|
COMPRESSED_VIDEO_FILE = 'compressed_video.avi'
|
|
MAX_FRAMES = 0 # Limit the number of frames processed
|
|
CRF = 0
|
|
SPEED = PRESET_SPEED_CATEGORIES.index("ultrafast")
|
|
|
|
# Load the trained model
|
|
MODEL = tf.keras.models.load_model('models/model.tf', custom_objects={'VideoCompressionModel': VideoCompressionModel, 'psnr': psnr})
|
|
|
|
# Load the uncompressed video
|
|
UNCOMPRESSED_VIDEO_FILE = 'test_data/B4_t02.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
|
|
cap.release()
|
|
|
|
return frame
|
|
|
|
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)
|
|
scaled_crf = scale_crf(CRF)
|
|
scaled_speed = scale_speed_preset(SPEED)
|
|
|
|
frame = combine_batch(uncompressed_frame, scaled_crf, scaled_speed, resize=False)
|
|
|
|
compressed_frame = MODEL.predict([np.expand_dims(frame, axis=0)])[0]
|
|
|
|
compressed_frame = compressed_frame[:, :, :3] # Keep only the first 3 channels (BGR)
|
|
|
|
compressed_frame = np.clip(compressed_frame * 255.0, 0, 255).astype(np.uint8)
|
|
|
|
cv2.imshow("comp", compressed_frame)
|
|
|
|
return compressed_frame
|
|
|
|
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), True)
|
|
|
|
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)
|
|
|
|
compressed_frame = cv2.resize(compressed_frame, (width, height))
|
|
|
|
#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)
|
|
|
|
#if i % 10 == 0: # Print progress every 10 frames
|
|
# print(f"Processed {i} / {total_frames} frames")
|
|
|
|
out.release()
|
|
print("Compression completed.")
|
|
|