jordon/DeepEncode
Archived
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

working

Browse source
This commit is contained in:
Jordon Brooks 2023-08-23 00:54:06 +01:00
parent f4512bba99
commit db43239b3d
5 changed files with 311 additions and 197 deletions
Download patch file Download diff file Expand all files Collapse all files

181
DeepEncode.py
View file

@ -1,90 +1,145 @@
# DeepEncode.py
import os
import argparse
import cv2
import numpy as np
from featureExtraction import combined, preprocess_frame, psnr, scale_crf, scale_speed_preset, ssim
from globalVars import PRESET_SPEED_CATEGORIES
# Set TensorFlow log level before any other imports
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
import tensorflow as tf
import numpy as np
import cv2
from featureExtraction import combined, combined_loss, psnr, scale_crf, scale_speed_preset, ssim
from globalVars import PRESET_SPEED_CATEGORIES, clear_screen
from video_compression_model import VideoCompressionModel, combine_batch
# Constants
COMPRESSED_VIDEO_FILE = 'compressed_video.avi'
MAX_FRAMES = 200 # Limit the number of frames processed
CRF = 51
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, 'ssim': ssim, 'combined': combined})
# Load the uncompressed video
CRF = 10
SPEED = "ultrafast"
MODEL_PATH = 'models/model.tf'
UNCOMPRESSED_VIDEO_FILE = 'test_data/x264_crf-5_preset-veryslow.mkv'
DISPLAY_OUTPUT = False
CROP_DIMENSIONS = None
def parse_arguments():
global COMPRESSED_VIDEO_FILE, MAX_FRAMES, CRF, SPEED, MODEL_PATH, UNCOMPRESSED_VIDEO_FILE, DISPLAY_OUTPUT, CROP_DIMENSIONS
parser = argparse.ArgumentParser(description='Deep Encoding of Videos')
parser.add_argument('-o', '--compressed_video_file', default=COMPRESSED_VIDEO_FILE, help='Path to the compressed video file')
parser.add_argument('-m', '--max_frames', type=int, default=MAX_FRAMES, help='Maximum number of frames to process')
parser.add_argument('-c', '--crf', type=int, default=CRF, help='CRF value for video compression')
parser.add_argument('-s', '--speed', default=SPEED, choices=PRESET_SPEED_CATEGORIES, help='Video compression speed category')
parser.add_argument('-p', '--model_path', default=MODEL_PATH, help='Path to the trained model')
parser.add_argument('-i', '--uncompressed_video_file', default=UNCOMPRESSED_VIDEO_FILE, help='Path to the uncompressed video file')
parser.add_argument('-d', '--display_output', action='store_true', default=DISPLAY_OUTPUT, help='Display real-time output to screen')
parser.add_argument('--keep_black_bars', action='store_true', help='Keep black bars from the video', default=False)
args = parser.parse_args()
COMPRESSED_VIDEO_FILE = args.compressed_video_file
MAX_FRAMES = args.max_frames
CRF = args.crf
SPEED = args.speed
MODEL_PATH = args.model_path
UNCOMPRESSED_VIDEO_FILE = args.uncompressed_video_file
DISPLAY_OUTPUT = args.display_output
if not args.keep_black_bars:
CROP_DIMENSIONS = find_crop_dimensions(UNCOMPRESSED_VIDEO_FILE)
def crop_black_bars(frame):
# Convert to grayscale for easier processing
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
# Threshold the image to make everything below a certain gray value black, and everything else white
_, thresh = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY)
# Find the contours of the white regions
contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# Find the bounding box that contains all the contours
x_min = y_min = float('inf')
x_max = y_max = 0
for contour in contours:
x, y, w, h = cv2.boundingRect(contour)
x_min = min(x_min, x)
y_min = min(y_min, y)
x_max = max(x_max, x + w)
y_max = max(y_max, y + h)
return x_min, y_min, x_max, y_max
def find_crop_dimensions(video_file):
cap = cv2.VideoCapture(video_file)
while True:
ret, frame = cap.read()
if not ret:
print("Error: Unable to find a non-black frame.")
cap.release()
exit()
# Check if the frame is entirely black
if np.any(frame > 0):
x_min, y_min, x_max, y_max = crop_black_bars(frame)
cap.release()
return x_min, y_min, x_max, y_max
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
return frame if ret else None
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)
def predict_frame(uncompressed_frame, model, crf, speed):
scaled_crf = scale_crf(crf)
scaled_speed = scale_speed_preset(PRESET_SPEED_CATEGORIES.index(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)
cv2.waitKey(1)
return compressed_frame
compressed_frame = model.predict([np.expand_dims(frame, axis=0)])[0]
return np.clip(compressed_frame[:, :, :3] * 255.0, 0, 255).astype(np.uint8)
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))
fps = int(cap.get(cv2.CAP_PROP_FPS))
cap.release()
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter(COMPRESSED_VIDEO_FILE, fourcc, fps, (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)
def main():
model = tf.keras.models.load_model(MODEL_PATH, custom_objects={'VideoCompressionModel': VideoCompressionModel, 'psnr': psnr, 'ssim': ssim, 'combined': combined, 'combined_loss': combined_loss})
cap = cv2.VideoCapture(UNCOMPRESSED_VIDEO_FILE)
compressed_frame = cv2.resize(compressed_frame, (width, height))
total_frames = min(int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), MAX_FRAMES)
height, width, fps = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FPS))
#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)
cap.release()
#if i % 10 == 0: # Print progress every 10 frames
# print(f"Processed {i} / {total_frames} frames")
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter(COMPRESSED_VIDEO_FILE, fourcc, fps, (width, height), True)
out.release()
print("Compression completed.")
if not out.isOpened():
print("Error: VideoWriter could not be opened.")
exit()
for i in range(total_frames):
uncompressed_frame = load_frame_from_video(UNCOMPRESSED_VIDEO_FILE, frame_num=i)
if CROP_DIMENSIONS:
x_min, y_min, x_max, y_max = CROP_DIMENSIONS
uncompressed_frame = uncompressed_frame[y_min:y_max, x_min:x_max]
compressed_frame = predict_frame(uncompressed_frame, model, CRF, SPEED)
compressed_frame = cv2.resize(compressed_frame, (width, height))
out.write(compressed_frame)
if DISPLAY_OUTPUT:
cv2.imshow('Compressed Video', compressed_frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
out.release()
print("Compression completed.")
if __name__ == '__main__':
clear_screen()
parse_arguments()
main()