video_utils.py

# -*- coding: utf-8 -*-
"""checker_c.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1FUjmVbfkq-N8h2yjbzl5eFez24jh0AV5
"""
import cv2
import torch
import torchvision.transforms as transforms
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from skimage.metrics import structural_similarity as ssim
from skimage.metrics import mean_squared_error as mse
import streamlit as st 
import torch.nn as nn


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')


def denormalize(tensor):
    # Assumes the input tensor was normalized with mean=0.5 and std=0.5
    tensor = tensor * 0.5 + 0.5
    tensor = torch.clamp(tensor, 0, 1)  # Ensures pixel values are between 0 and 1
    return tensor


# Helper functions for preprocessing, sharpening, and denormalizing
def preprocess_frames(frames, preprocess):
    frames_tensor = []
    for frame in frames:
        frame_tensor = preprocess(frame)
        frames_tensor.append(frame_tensor)
    return torch.stack(frames_tensor)

def sharpen_image(image):
    kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]], dtype=np.float32)
    return cv2.filter2D(image, -1, kernel)

# Function to resize tensor images
def resize_tensor(tensor, size, interpolation=Image.BILINEAR):
    transform = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((size[1], size[0]), interpolation=interpolation),  # Correct order for PIL
        transforms.ToTensor()
    ])
    return transform(tensor)


# Calculate MSE and SSIM
def calculate_metrics(actual_frames, predicted_frames):

    mse_values = []
    ssim_values = []
    
    for actual, predicted in zip(actual_frames, predicted_frames):
        # Convert to numpy and transpose to (H, W, channels)
        actual_np = actual.permute(1, 2, 0).numpy()
        predicted_np = predicted.permute(1, 2, 0).numpy()
        
        # Ensure values are in range [0, 1] for SSIM
        actual_np = np.clip(actual_np, 0, 1)
        predicted_np = np.clip(predicted_np, 0, 1)
        
        # Convert to grayscale for SSIM if multichannel
        if actual_np.shape[2] > 1:
            actual_gray = np.dot(actual_np[..., :3], [0.2989, 0.5870, 0.1140])
            predicted_gray = np.dot(predicted_np[..., :3], [0.2989, 0.5870, 0.1140])
        else:
            actual_gray = actual_np.squeeze()
            predicted_gray = predicted_np.squeeze()
        
        # Calculate MSE
        mse = np.mean((actual_np - predicted_np) ** 2)
        mse_values.append(mse)
        
        # Calculate SSIM
        ssim_value = ssim(actual_gray, predicted_gray, data_range=1.0)
        ssim_values.append(ssim_value)
    
    return mse_values, ssim_values


# Function to read video frames
def predict_video_sequences(model, video_path, num_sequences, device, 
                             input_frames=10, output_frames=5, 
                             frame_size=(64,64)):

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Error opening video file {video_path}")
        return None

    total_frames_in_video = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    preprocess = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((frame_size[1], frame_size[0])),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5]*3, std=[0.5]*3)
    ])

    input_sequences = []
    predicted_sequences = []
    actual_sequences = []
    metrics_sequences = []

    sequences_processed = 0
    frame_idx = 0

    while sequences_processed < num_sequences and frame_idx + input_frames <= total_frames_in_video:
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)

        # Read input and output frames
        input_frame_list, actual_output_frames = read_video_frames(
            cap, input_frames, output_frames)

        if len(input_frame_list) < input_frames or len(actual_output_frames) < output_frames:
            break

        # Preprocess input and output frames
        input_frames_tensor = preprocess_frames(input_frame_list, preprocess)
        actual_output_frames_tensor = preprocess_frames(actual_output_frames, preprocess)

        # Predict frames
        model.eval()
        with torch.no_grad():
            predicted_frames = model(input_frames_tensor.unsqueeze(0).to(device))
        predicted_frames = predicted_frames.squeeze(0).cpu()

        # Calculate metrics
        mse_values, ssim_values = calculate_metrics(actual_output_frames_tensor, predicted_frames)

        # Store sequences and metrics
        input_sequences.append(input_frame_list)
        predicted_sequences.append(predicted_frames.numpy())
        actual_sequences.append(actual_output_frames)
        metrics_sequences.append({'mse': mse_values, 'ssim': ssim_values})

        sequences_processed += 1
        frame_idx += input_frames + output_frames

    cap.release()
    return input_sequences, predicted_sequences, actual_sequences, metrics_sequences

# Function to generate actual video
def generate_actual_video(video_path, output_path_actual, frame_size=(320, 240), frame_rate=30):
  
    # Define the preprocessing pipeline
    preprocess = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((64, 64)),  # Assuming the model expects 64x64 inputs
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5]*3, std=[0.5]*3)
    ])

    # Initialize video capture
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Error opening video file {video_path}")
        return

    total_frames_in_video = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    print(f"Total frames in video: {total_frames_in_video}")

    # Initialize video writer for the actual video
    fourcc = cv2.VideoWriter_fourcc(*'avc1')
    out_actual = cv2.VideoWriter(output_path_actual, fourcc, frame_rate, (frame_size[0], frame_size[1]))

    all_frames = []
    frame_count = 0

    # Read all frames from the video
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        all_frames.append(frame_rgb)
        frame_count += 1
        if frame_count % 100 == 0:
            print(f"Read {frame_count}/{total_frames_in_video} frames")

    cap.release()
    print(f"Finished reading {frame_count} frames from the video.")

    # Preprocess all frames
    frames_tensor = preprocess_frames(all_frames, preprocess)
    print("Preprocessing completed.")

    # Denormalize the frames
    frames_denorm = denormalize(frames_tensor)
    print("Denormalization completed.")

    # Resize all frames to the desired frame size
    frames_resized = [resize_tensor(frame, frame_size) for frame in frames_denorm]
    print("Resizing completed.")

    # Write all resized frames to the actual video
    for idx, frame in enumerate(frames_resized):
        frame_np = frame.permute(1, 2, 0).numpy()  # Convert from (C, H, W) to (H, W, C)
        frame_np = np.clip(frame_np * 255.0, 0, 255).astype(np.uint8)  # Convert to uint8
        frame_bgr = cv2.cvtColor(frame_np, cv2.COLOR_RGB2BGR)  # Convert back to BGR for OpenCV
        out_actual.write(frame_bgr)
        if (idx + 1) % 100 == 0:
            print(f"Wrote {idx + 1}/{len(frames_resized)} frames to the actual video.")

    out_actual.release()
    print(f'Actual video saved to {output_path_actual}')

# Function to display video sequences
def display_video_sequences(input_sequences, predicted_sequences, 
                            metrics_sequences, 
                            display_size=(320, 240), 
                            is_sharpening=True):
    """
    Display video sequences with input and predicted frames.
    
    Args:
        input_sequences: List of input frame sequences (each sequence is a list of frames)
        predicted_sequences: List of predicted frame sequences (each sequence is a list of frames)
        metrics_sequences: List of metric dictionaries (each dict contains 'mse' and 'ssim' lists)
        display_size: Size to display frames as (width, height)
        is_sharpening: Whether to apply image sharpening to predicted frames
    """
    for seq_idx, (inputs, predicts, metrics) in enumerate(
        zip(input_sequences, predicted_sequences, metrics_sequences), 1):
        
        st.write(f"### Sequence {seq_idx} Metrics:")
        for i, (mse, ssim) in enumerate(zip(metrics['mse'], metrics['ssim']), 1):
            st.write(f"**Frame {i}** - MSE: {mse:.4f}, SSIM: {ssim:.4f}")

        # Determine the number of frames to display
        num_input = len(inputs)
        num_pred = len(predicts)
        num_cols = max(num_input, num_pred)

        # Create a figure with 2 rows and 'num_cols' columns
        fig, axs = plt.subplots(2, num_cols, figsize=(3 * num_cols, 6))  # Adjusted height for two rows

        # If there's only one column, make axs 2D for consistency
        if num_cols == 1:
            axs = axs[:, np.newaxis]

        # ============================
        # Display Input Frames (Row 1)
        # ============================
        for i, frame in enumerate(inputs):
            # Resize frame
            frame_resized = cv2.resize(frame, display_size, interpolation=cv2.INTER_LANCZOS4)
            
            # Normalize frame to [0,1] if it's in [0,255]
            if frame_resized.max() > 1.0:
                frame_resized = frame_resized / 255.0
            
            # Display frame without color conversion
            axs[0, i].imshow(frame_resized)
            axs[0, i].set_title(f'Input {i+1}')
            axs[0, i].axis('off')

        # Fill remaining input frames with empty plots if any
        for i in range(num_input, num_cols):
            axs[0, i].axis('off')

        # ==============================
        # Display Predicted Frames (Row 2)
        # ==============================
        for i, frame_tensor in enumerate(predicts):
            # Convert tensor to numpy array
            if isinstance(frame_tensor, torch.Tensor):
                frame_np = frame_tensor.cpu().detach().numpy()
            else:
                frame_np = frame_tensor

            # Handle channels: assume (C, H, W) format
            if frame_np.shape[0] == 3:
                frame_np = np.transpose(frame_np, (1, 2, 0))  # (H, W, C)
            elif frame_np.shape[-1] != 3:
                st.warning(f"Predicted frame {i+1} has unexpected shape: {frame_np.shape}")
                axs[1, i].axis('off')
                continue  # Skip frames with unexpected shapes

            # Denormalize if necessary (adjust based on your model's normalization)
            # Example: If your model normalized with mean=0.5 and std=0.5
            frame_denorm = frame_np * 0.5 + 0.5  # [0,1] range
            frame_denorm = np.clip(frame_denorm, 0, 1)

            # Resize frame
            frame_resized = cv2.resize(frame_denorm, display_size, interpolation=cv2.INTER_LANCZOS4)
            
            # Apply sharpening if enabled
            if is_sharpening:
                frame_resized = sharpen_image(frame_resized)
                frame_resized = np.clip(frame_resized, 0, 1)  # Ensure values are within [0,1]

            # Display predicted frame
            axs[1, i].imshow(frame_resized)
            axs[1, i].set_title(f'Predicted {i+1}')
            axs[1, i].axis('off')

        # Fill remaining predicted frames with empty plots if any
        for i in range(num_pred, num_cols):
            axs[1, i].axis('off')

        plt.tight_layout()
        st.pyplot(fig)



def generate_videos(model, video_path, output_path_predicted, 
                    frame_size=(320, 240), frame_rate=30, 
                    input_frames=10, output_frames=5):

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    model.eval()

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Error opening video file {video_path}")
        return

    total_frames_in_video = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    # Initialize video writers
    fourcc = cv2.VideoWriter_fourcc(*'avc1')
    out_predicted = cv2.VideoWriter(output_path_predicted, fourcc, frame_rate, (frame_size[0], frame_size[1]))
    
    preprocess = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((64, 64)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5]*3, std=[0.5]*3)
    ])

    while True:
        # Read input frames
        input_frame_list = []
        for _ in range(input_frames):
            ret, frame = cap.read()
            if not ret:
                print("End of video reached or cannot read the frame.")
                break
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            input_frame_list.append(frame_rgb)

        if len(input_frame_list) < input_frames:
            break

        # Preprocess and predict
        input_frames_tensor = preprocess_frames(input_frame_list, preprocess)
        with torch.no_grad():
            predicted_frames = model(input_frames_tensor.unsqueeze(0).to(device))
        predicted_frames = predicted_frames.squeeze(0).cpu()

        # Prepare input and predicted video frames
        input_frames_denorm = denormalize(input_frames_tensor)
        predicted_frames_denorm = denormalize(predicted_frames)

        input_frames_resized = [resize_tensor(frame, frame_size) for frame in input_frames_denorm]
        predicted_frames_resized = [resize_tensor(frame, frame_size) for frame in predicted_frames_denorm]

        # Write input and predicted frames to predicted video
        for frame in input_frames_resized + predicted_frames_resized:
            frame_np = frame.permute(1, 2, 0).numpy()
            frame_np = np.clip(frame_np * 255.0, 0, 255).astype(np.uint8)
            frame_bgr = cv2.cvtColor(frame_np, cv2.COLOR_RGB2BGR)
            out_predicted.write(frame_bgr)

        # Write actual frames to actual video
        actual_frames_resized = []
        for _ in range(output_frames):
            ret, frame = cap.read()
            if not ret:
                break
            frame_resized = cv2.resize(frame, frame_size)
            actual_frames_resized.append(frame_resized)

        # Check for end of video
        pos_frames = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
        if pos_frames + input_frames > total_frames_in_video:
            break

    # Release resources
    cap.release()
    out_predicted.release()
    print(f'Predicted video saved to {output_path_predicted}')


# Utility functions (these would be in a separate utility module in a real project)
def read_video_frames(cap, input_frames, output_frames):
    input_frame_list = []
    actual_output_frames = []
    
    for i in range(input_frames + output_frames):
        ret, frame = cap.read()
        if not ret:
            break
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        if i < input_frames:
            input_frame_list.append(frame_rgb)
        else:
            actual_output_frames.append(frame_rgb)
    
    return input_frame_list, actual_output_frames