Initial Commit

DeepEncode.py

@@ -0,0 +1,64 @@
+import tensorflow as tf
+import numpy as np
+import cv2
+from video_compression_model import VideoCompressionModel
+
+# Constants
+NUM_CHANNELS = 3
+
+# Step 2: Load the trained model
+model = tf.keras.models.load_model('ai_rate_control_model.keras', custom_objects={'VideoCompressionModel': VideoCompressionModel})
+
+# Step 3: Load the uncompressed video
+UNCOMPRESSED_VIDEO_FILE = 'test_video.mkv'
+
+def load_frames_from_video(video_file, num_frames = 0):
+    print("Extracting video frames...")
+    cap = cv2.VideoCapture(video_file)
+    frames = []
+    count = 0
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            print("Max frames from file reached")
+            break
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frames.append(frame)
+        count += 1
+        if num_frames == 0 or count >= num_frames:
+            print("Max Frames wanted reached: ", num_frames)
+            break
+    cap.release()
+    print("Extraction Complete")
+    return frames
+
+uncompressed_frames = load_frames_from_video(UNCOMPRESSED_VIDEO_FILE, 200)
+if len(uncompressed_frames) == 0 or None:
+    print("IO ERROR!")
+    exit()
+    
+uncompressed_frames = np.array(uncompressed_frames) / 255.0
+
+if len(uncompressed_frames) == 0 or None:
+    print("np.array ERROR!")
+    exit()
+
+# Step 4: Compress the video frames using the loaded model
+compressed_frames = model.predict(uncompressed_frames)
+
+# Step 5: Save the compressed video frames
+COMPRESSED_VIDEO_FILE = 'compressed_video.mkv'
+
+def save_frames_as_video(frames, video_file):
+    print("Saving video frames...")
+    height, width = frames[0].shape[:2]
+    fourcc = cv2.VideoWriter_fourcc(*'XVID')
+    out = cv2.VideoWriter(video_file, fourcc, 24.0, (width, height))
+    for frame in frames:
+        frame = np.clip(frame * 255.0, 0, 255).astype(np.uint8)
+        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+        out.write(frame)
+    out.release()
+
+save_frames_as_video(compressed_frames, COMPRESSED_VIDEO_FILE)
+print("Compression completed.")