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

update

Browse source
This commit is contained in:
Jordon Brooks 2023-08-17 01:57:53 +01:00
parent 3ea1568ad3
commit 7787d0584e
4 changed files with 35 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

13
DeepEncode.py
View file

@ -22,7 +22,7 @@ SPEED = PRESET_SPEED_CATEGORIES.index("ultrafast")
MODEL = tf.keras.models.load_model('models/model.tf', custom_objects={'VideoCompressionModel': VideoCompressionModel, 'psnr': psnr})
# Load the uncompressed video
UNCOMPRESSED_VIDEO_FILE = 'test_data/training_video.mkv'
UNCOMPRESSED_VIDEO_FILE = 'test_data/B4_t02.mkv'
def load_frame_from_video(video_file, frame_num):
cap = cv2.VideoCapture(video_file)
@ -41,16 +41,15 @@ def predict_frame(uncompressed_frame):
scaled_crf = scale_crf(CRF)
scaled_speed = scale_speed_preset(SPEED)
frame = combine_batch(uncompressed_frame, scaled_crf, scaled_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)
display_frame = np.clip(compressed_frame * 255.0, 0, 255).astype(np.uint8)
compressed_frame = np.clip(compressed_frame * 255.0, 0, 255).astype(np.uint8)
cv2.imshow("comp", display_frame)
cv2.waitKey(1)
cv2.imshow("comp", compressed_frame)
return compressed_frame
@ -75,9 +74,9 @@ for i in range(total_frames):
compressed_frame = cv2.resize(compressed_frame, (width, height))
compressed_frame = np.clip(compressed_frame * 255.0, 0, 255).astype(np.uint8)
#compressed_frame = np.clip(compressed_frame * 255.0, 0, 255).astype(np.uint8)
compressed_frame = cv2.cvtColor(compressed_frame, cv2.COLOR_RGB2BGR)
#compressed_frame = cv2.cvtColor(compressed_frame, cv2.COLOR_RGB2BGR)
out.write(compressed_frame)