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

Working GPU model

Browse source
This commit is contained in:
Jordon Brooks 2023-07-30 11:49:19 +01:00
parent 5085c87300
commit dea59068fb
3 changed files with 190 additions and 108 deletions
Download patch file Download diff file Expand all files Collapse all files

52
DeepEncode.py
View file

@ -1,15 +1,19 @@
# DeepEncode.py
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
CHUNK_SIZE = 10 # Adjust based on available memory and video resolution
COMPRESSED_VIDEO_FILE = 'compressed_video.avi'
MAX_FRAMES = 0 # Limit the number of frames processed
CRF = 25.0 # Example CRF value
PRESET_SPEED = 4 # Index for "fast" in our defined list
# Load the trained model
model = tf.keras.models.load_model('models/model.keras', custom_objects={'VideoCompressionModel': VideoCompressionModel})
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'
@ -23,46 +27,56 @@ def load_frame_from_video(video_file, frame_num):
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])
# 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(10)
compressed_frame = model.predict({
"frame": np.array([uncompressed_frame]),
"uncompressed_frame": uncompressed_frame,
"compressed_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(10)
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
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 = predict_frame(uncompressed_frame, model, CRF, 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)
#cv2.imshow("output", compressed_frame)
out.release()
print("Compression completed.")