DeepEncode/train_model.py

import os
import json
import tensorflow as tf
import numpy as np
import cv2
from video_compression_model import NUM_CHANNELS, NUM_FRAMES, VideoCompressionModel, PRESET_SPEED_CATEGORIES
from tensorflow.keras.callbacks import EarlyStopping

print(tf.config.list_physical_devices('GPU'))

# Constants
BATCH_SIZE = 8
EPOCHS = 5
TRAIN_SAMPLES = 5

def load_list(list_path):
    with open(list_path, "r") as json_file:
        video_details_list = json.load(json_file)
    return video_details_list

def load_frames_from_video(video_file, num_frames):
    print("Extracting video frames...")
    cap = cv2.VideoCapture(video_file)
    frames = []
    count = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frames.append(frame)
        count += 1
        if count >= num_frames:
            break
    cap.release()
    width, height = frame.shape[:2]
    return frames, width, height

def preprocess(frames):
    return np.array(frames) / 255.0

def save_model(model, file):
    os.makedirs("models", exist_ok=True)
    model.save(os.path.join("models/", file))
    print("Model saved successfully!")

def load_video_from_list(list_path):
    details_list = load_list(list_path)
    all_frames = []
    all_details = []
    for video_details in details_list:
        VIDEO_FILE = video_details["video_file"]
        CRF = video_details['crf'] / 63.0
        PRESET_SPEED = PRESET_SPEED_CATEGORIES.index(video_details['preset_speed'])
        video_details['preset_speed'] = PRESET_SPEED

        train_frames, w, h = load_frames_from_video(os.path.join("test_data/", VIDEO_FILE), NUM_FRAMES * TRAIN_SAMPLES)
        
        all_frames.extend(train_frames)
        all_details.append({
            "frames": train_frames,
            "width": w,
            "height": h,
            "crf": CRF,
            "preset_speed": PRESET_SPEED,
            "video_file": VIDEO_FILE
        })
    return all_details

def generate_frame_sequences(frames):
    sequences = []
    labels = []
    for i in range(len(frames) - NUM_FRAMES + 1):
        sequence = frames[i:i+NUM_FRAMES-1]
        sequences.append(sequence)
        labels.append(sequence[-1])
    return np.array(sequences), np.array(labels)

def frame_difference(frames):
    differences = []
    for i in range(1, len(frames)):
        differences.append(cv2.absdiff(frames[i], frames[i-1]))
    return differences

def main():
    all_video_details_train = load_video_from_list("test_data/training.json")
    all_video_details_val = load_video_from_list("test_data/validation.json")

    model = VideoCompressionModel(NUM_CHANNELS, NUM_FRAMES)
    model.compile(loss='mean_squared_error', optimizer='adam')
    early_stop = EarlyStopping(monitor='val_loss', patience=3, verbose=1, restore_best_weights=True)

    # Load and concatenate all sequences and labels
    all_train_sequences = []
    all_val_sequences = []
    all_train_labels = []
    all_val_labels = []
    all_crf_train = []
    all_crf_val = []
    all_preset_speed_train = []
    all_preset_speed_val = []
    
    for video_details_train, video_details_val in zip(all_video_details_train, all_video_details_val):
        train_frames = video_details_train["frames"]
        val_frames = video_details_val["frames"]

        train_differences = frame_difference(preprocess(train_frames))
        val_differences = frame_difference(preprocess(val_frames))
        
        #print(len(train_differences), train_differences[0].shape)

        train_sequences, train_labels = generate_frame_sequences(train_differences)
        val_sequences, val_labels = generate_frame_sequences(val_differences)

        crf_array_train = np.full((len(train_sequences), 1), video_details_train['crf'])
        crf_array_val = np.full((len(val_sequences), 1), video_details_val['crf'])
        preset_speed_array_train = np.full((len(train_sequences), 1), video_details_train['preset_speed'])
        preset_speed_array_val = np.full((len(val_sequences), 1), video_details_val['preset_speed'])

        all_train_sequences.extend(train_sequences)
        all_val_sequences.extend(val_sequences)
        all_train_labels.extend(train_labels)
        all_val_labels.extend(val_labels)
        all_crf_train.extend(crf_array_train)
        all_crf_val.extend(crf_array_val)
        all_preset_speed_train.extend(preset_speed_array_train)
        all_preset_speed_val.extend(preset_speed_array_val)

    # Convert lists to numpy arrays
    all_train_sequences = np.array(all_train_sequences)
    all_val_sequences = np.array(all_val_sequences)
    all_train_labels = np.array(all_train_labels)
    all_val_labels = np.array(all_val_labels)
    all_crf_train = np.array(all_crf_train)
    all_crf_val = np.array(all_crf_val)
    all_preset_speed_train = np.array(all_preset_speed_train)
    all_preset_speed_val = np.array(all_preset_speed_val)

    # Shuffle the training data
    indices_train = np.arange(all_train_sequences.shape[0])
    np.random.shuffle(indices_train)

    all_train_sequences = all_train_sequences[indices_train]
    all_train_labels = all_train_labels[indices_train]
    all_crf_train = all_crf_train[indices_train]
    all_preset_speed_train = all_preset_speed_train[indices_train]

    print("\nTraining the model on mixed sequences...")
    model.fit(
        {"frames": all_train_sequences, "crf": all_crf_train, "preset_speed": all_preset_speed_train},
        all_train_labels,
        batch_size=BATCH_SIZE,
        epochs=EPOCHS,
        validation_data=({"frames": all_val_sequences, "crf": all_crf_val, "preset_speed": all_preset_speed_val}, all_val_labels),
        callbacks=[early_stop]
    )
    print("\nTraining completed!")

    save_model(model, 'model_differencing.keras')

if __name__ == "__main__":
    main()