optimisation

DeepEncode.py

@@ -4,82 +4,88 @@ import cv2
 from video_compression_model import NUM_FRAMES, PRESET_SPEED_CATEGORIES, VideoCompressionModel
 
 # Constants
-NUM_CHANNELS = 3
+MAX_FRAMES = 24
+CHUNK_SIZE = 24  # Adjust based on available memory and video resolution
+COMPRESSED_VIDEO_FILE = 'compressed_video.mkv'
+
 
 # Step 2: 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_differencing.keras', custom_objects={'VideoCompressionModel': VideoCompressionModel})
 
 # Step 3: Load the uncompressed video
 UNCOMPRESSED_VIDEO_FILE = 'test_data/training_video.mkv'
 
-def load_frames_from_video(video_file, num_frames = 0):
-    print("Extracting video frames...")
+def load_frames_from_video(video_file, start_frame=0, num_frames=CHUNK_SIZE):
     cap = cv2.VideoCapture(video_file)
     frames = []
-    count = 0
-    while True:
+    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
+    
+    for _ in range(num_frames):
         ret, frame = cap.read()
         if not ret:
-            print("Max frames from file reached")
             break
-        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0  # Normalize and convert to float32
         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, 100)
-if not uncompressed_frames:
-    print("IO ERROR!")
-    exit()
+def predict_in_chunks(uncompressed_frames, model, crf_values, preset_speed_values):
+    num_sequences = len(uncompressed_frames) - NUM_FRAMES + 1
+    compressed_frames = []
     
-uncompressed_frames = np.array(uncompressed_frames) / 255.0
+    for frame in uncompressed_frames:
+        cv2.imshow("frame", frame)
+        cv2.waitKey(50)
 
-# Generate sequences of frames for prediction
-uncompressed_frame_sequences = []
-for i in range(len(uncompressed_frames) - NUM_FRAMES + 1):
-    sequence = uncompressed_frames[i:i+NUM_FRAMES]
-    uncompressed_frame_sequences.append(sequence)
-uncompressed_frame_sequences = np.array(uncompressed_frame_sequences)
+    for start in range(0, num_sequences, CHUNK_SIZE):
+        end = min(start + CHUNK_SIZE, num_sequences)
+        frame_chunk = uncompressed_frames[start:end + NUM_FRAMES - 1]
+        crf_chunk = crf_values[start:end]
+        speed_chunk = preset_speed_values[start:end]
 
-#for frame in uncompressed_frames:
-#    cv2.imshow('Frame', frame)
-#    cv2.waitKey(50)  # Display each frame for 1 second
+        frame_sequences = []
+        for i in range(len(frame_chunk) - NUM_FRAMES + 1):
+            sequence = frame_chunk[i:i + NUM_FRAMES]
+            frame_sequences.append(sequence)
+            
+        frame_sequences = np.array(frame_sequences)
 
+        compressed_chunk = model.predict({"frames": frame_sequences, "crf": crf_chunk, "preset_speed": speed_chunk})
+        compressed_frames.extend(compressed_chunk)
+        
+    return compressed_frames
 
-# Step 4: Compress the video frames using the loaded model
-crf_values = np.full((len(uncompressed_frame_sequences), 1), 25, dtype=np.float32)  # Added dtype argument
-
-preset_speed_index = PRESET_SPEED_CATEGORIES.index("fast")
-preset_speed_values = np.full((len(uncompressed_frame_sequences), 1), preset_speed_index, dtype=np.float32)  # Added dtype argument
-
-compressed_frame_sequences = model.predict({"frames": uncompressed_frame_sequences, "crf": crf_values, "preset_speed": preset_speed_values})
-
-# We'll use the last frame of each sequence as the compressed frame
-#compressed_frames = compressed_frame_sequences[:, -1]
-
-#for frame in compressed_frame_sequences:
-#    cv2.imshow('Compressed Frame', frame)
-#    cv2.waitKey(50)
-
-
-# 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))
+def save_frames_chunk(frames, video_writer):
     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()
+        video_writer.write(frame)
 
-save_frames_as_video(compressed_frame_sequences, COMPRESSED_VIDEO_FILE)
+
+cap = cv2.VideoCapture(UNCOMPRESSED_VIDEO_FILE)
+total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+if MAX_FRAMES != 0 and total_frames > MAX_FRAMES:
+    total_frames = MAX_FRAMES
+
+cap.release()
+
+crf_values = np.full((CHUNK_SIZE + NUM_FRAMES - 1, 1), 25, dtype=np.float32)  # Chunk size + look-ahead frames
+preset_speed_index = PRESET_SPEED_CATEGORIES.index("fast")
+preset_speed_values = np.full((CHUNK_SIZE + NUM_FRAMES - 1, 1), preset_speed_index, dtype=np.float32)
+
+out = None  # Video writer instance
+for i in range(0, total_frames, CHUNK_SIZE):
+    uncompressed_frames_chunk = load_frames_from_video(UNCOMPRESSED_VIDEO_FILE, start_frame=i)
+    compressed_frames_chunk = predict_in_chunks(uncompressed_frames_chunk, model, crf_values, preset_speed_values)
+    
+    # Initialize video writer if it's the first chunk
+    if out is None:
+        height, width = compressed_frames_chunk[0].shape[:2]
+        fourcc = cv2.VideoWriter_fourcc(*'XVID')
+        out = cv2.VideoWriter(COMPRESSED_VIDEO_FILE, fourcc, 24.0, (width, height))
+    
+    save_frames_chunk(compressed_frames_chunk, out)
+
+out.release()
 print("Compression completed.")