140 lines
No EOL
4.8 KiB
Python
140 lines
No EOL
4.8 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import cv2
|
|
import numpy as np
|
|
|
|
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
|
|
|
|
import tensorflow as tf
|
|
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
|
|
|
|
from video_compression_model import VideoCompressionModel, data_generator
|
|
|
|
from globalVars import HEIGHT, WIDTH, LOGGER
|
|
|
|
# Constants
|
|
BATCH_SIZE = 16
|
|
EPOCHS = 5
|
|
LEARNING_RATE = 0.01
|
|
MODEL_SAVE_FILE = "models/model.tf"
|
|
MODEL_CHECKPOINT_DIR = "checkpoints"
|
|
EARLY_STOP = 10
|
|
|
|
def save_model(model):
|
|
try:
|
|
LOGGER.debug("Attempting to save the model.")
|
|
os.makedirs("models", exist_ok=True)
|
|
model.save(MODEL_SAVE_FILE, save_format='tf')
|
|
LOGGER.info("Model saved successfully!")
|
|
except Exception as e:
|
|
LOGGER.error(f"Error saving the model: {e}")
|
|
raise
|
|
|
|
|
|
def load_video_metadata(list_path):
|
|
"""
|
|
Load video metadata from a JSON file.
|
|
|
|
Args:
|
|
- json_path (str): Path to the JSON file containing video metadata.
|
|
|
|
Returns:
|
|
- list: List of dictionaries, each containing video details.
|
|
"""
|
|
|
|
LOGGER.trace(f"Entering: load_video_metadata({list_path})")
|
|
try:
|
|
with open(list_path, "r") as json_file:
|
|
file = json.load(json_file)
|
|
LOGGER.trace(f"load_video_metadata returning: {file}")
|
|
return file
|
|
except FileNotFoundError:
|
|
LOGGER.error(f"Metadata file {list_path} not found.")
|
|
raise
|
|
except json.JSONDecodeError:
|
|
LOGGER.error(f"Error decoding JSON from {list_path}.")
|
|
raise
|
|
|
|
|
|
def main():
|
|
global BATCH_SIZE, EPOCHS, LEARNING_RATE, MODEL_SAVE_FILE
|
|
# Argument parsing
|
|
parser = argparse.ArgumentParser(description="Train the video compression model.")
|
|
parser.add_argument('-b', '--batch_size', type=int, default=BATCH_SIZE, help='Batch size for training.')
|
|
parser.add_argument('-e', '--epochs', type=int, default=EPOCHS, help='Number of epochs for training.')
|
|
parser.add_argument('-l', '--learning_rate', type=float, default=LEARNING_RATE, help='Learning rate for training.')
|
|
parser.add_argument('-c', '--continue_training', type=str, nargs='?', const=MODEL_SAVE_FILE, default=None, help='Path to the saved model to continue training. If used without a value, defaults to the MODEL_SAVE_FILE.')
|
|
args = parser.parse_args()
|
|
|
|
BATCH_SIZE = args.batch_size
|
|
EPOCHS = args.epochs
|
|
LEARNING_RATE = args.learning_rate
|
|
|
|
# Display training configuration
|
|
LOGGER.info("Starting the training with the given configuration.")
|
|
LOGGER.info("Training configuration:")
|
|
LOGGER.info(f"Batch size: {BATCH_SIZE}")
|
|
LOGGER.info(f"Epochs: {EPOCHS}")
|
|
LOGGER.info(f"Learning rate: {LEARNING_RATE}")
|
|
LOGGER.info(f"Continue training from: {MODEL_SAVE_FILE}")
|
|
|
|
LOGGER.debug(f"Max video resolution: {WIDTH}x{HEIGHT}")
|
|
LOGGER.trace("Hello, World!")
|
|
|
|
# Load all video metadata
|
|
all_videos = load_video_metadata("test_data/validation/validation.json")
|
|
|
|
# Split into training and validation
|
|
split_index = int(0.8 * len(all_videos))
|
|
training_videos = all_videos[:split_index]
|
|
validation_videos = all_videos[split_index:]
|
|
|
|
if args.continue_training:
|
|
model = tf.keras.models.load_model(args.continue_training)
|
|
else:
|
|
model = VideoCompressionModel()
|
|
|
|
|
|
# Set optimizer and compile the model
|
|
optimizer = tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE)
|
|
model.compile(loss='mean_squared_error', optimizer=optimizer)
|
|
|
|
# Define checkpoints and early stopping
|
|
checkpoint_callback = ModelCheckpoint(
|
|
filepath=os.path.join(MODEL_CHECKPOINT_DIR, "epoch-{epoch:02d}.tf"),
|
|
save_weights_only=False,
|
|
save_best_only=False,
|
|
verbose=1,
|
|
save_format="tf"
|
|
)
|
|
early_stop = EarlyStopping(monitor='val_loss', patience=EARLY_STOP, verbose=1, restore_best_weights=True)
|
|
|
|
# Calculate steps per epoch for training and validation
|
|
average_frames_per_video = 2880 # Given 2 minutes @ 24 fps
|
|
total_frames_train = average_frames_per_video * len(training_videos)
|
|
total_frames_validation = average_frames_per_video * len(validation_videos)
|
|
steps_per_epoch_train = total_frames_train // BATCH_SIZE
|
|
steps_per_epoch_validation = total_frames_validation // BATCH_SIZE
|
|
|
|
# Train the model
|
|
LOGGER.info("Starting model training.")
|
|
model.fit(
|
|
data_generator(training_videos, BATCH_SIZE),
|
|
epochs=EPOCHS,
|
|
steps_per_epoch=steps_per_epoch_train,
|
|
validation_data=data_generator(validation_videos, BATCH_SIZE), # Add validation data here
|
|
validation_steps=steps_per_epoch_validation, # Add validation steps here
|
|
callbacks=[early_stop, checkpoint_callback]
|
|
)
|
|
LOGGER.info("Model training completed.")
|
|
|
|
save_model(model)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except Exception as e:
|
|
LOGGER.error(f"Unexpected error during training: {e}")
|
|
raise |