quantize an entire video

quantize_mp4.py

import logging
from moviepy.editor import VideoFileClip, VideoClip, concatenate_videoclips
from PIL import Image
import numpy as np
from sklearn.cluster import KMeans
import random
import os

# Configure logging
logging.basicConfig(filename='script_log.log', level=logging.DEBUG)

# Function to apply k-means clustering to an image
def apply_kmeans(image, colors):
    im_array = np.array(image)
    reshaped_array = im_array.reshape((-1, 3))

    # Apply k-means clustering
    kmeans = KMeans(n_clusters=len(colors), random_state=42)
    labels = kmeans.fit_predict(reshaped_array)

    # Convert labels to integer array
    labels_int = labels.astype(int)

    # Assign the cluster centers to the image pixels
    segmented_image_array = np.array(colors)[labels_int].reshape(im_array.shape)

    # Convert the numpy array back to an image
    segmented_image = Image.fromarray(np.uint8(segmented_image_array))

    return segmented_image

# Define the colors and paths
def ColA():
    Rr = random.randint(0, 255)
    Gg = random.randint(0, 255)
    Bb = random.randint(0, 255)
    return [Rr, Gg, Bb]

colors = [
    ColA(),
    ColA(),
    ColA(),
    ColA(),    
    ColA(),
    ColA(),
    ColA(),
    ColA(),
    ColA()
]


input_video_path = "/home/jack/Desktop/vids/6bd83e3c-2f9f-49bc-ac9f-ee233870b241.mp4"
output_video_path = "junk/video_segmented.mp4"

# Create temporary directory
temp_directory = "temp_directory"
os.makedirs(temp_directory, exist_ok=True)

# Process each frame of the video
clip = VideoFileClip(input_video_path)
segmented_frames = []

# Apply k-means clustering to the entire video using the same set of colors
# Perform clustering only once and use the same colors for all frames
kmeans = KMeans(n_clusters=len(colors), random_state=42)
for i, frame in enumerate(clip.iter_frames(fps=clip.fps)):
    if i == 0:
        kmeans.fit(np.array(frame).reshape((-1, 3)))

    labels = kmeans.predict(np.array(frame).reshape((-1, 3)))
    segmented_frame_array = np.array(colors)[labels].reshape(np.array(frame).shape)
    segmented_frame = Image.fromarray(np.uint8(segmented_frame_array))
    segmented_frames.append(np.array(segmented_frame))

# Define a function to generate frames for the video
def make_frame(t):
    return segmented_frames[int(t * clip.fps)]

# Create a video clip from the segmented frames
logging.debug(f"Creating segmented video with {len(segmented_frames)} frames.")
segmented_clip = VideoClip(make_frame, duration=clip.duration)

# Write the segmented video to a file
segmented_clip.write_videofile(output_video_path, codec="libx264", audio_codec="aac", fps=clip.fps)

# Clean up temporary directory
for file_path in os.listdir(temp_directory):
    file_path = os.path.join(temp_directory, file_path)
    os.remove(file_path)
os.rmdir(temp_directory)

logging.info("Script execution completed.")

This script uses the MoviePy library to apply k-means clustering to each frame of a video, creating a segmented video where each pixel is assigned a color from a predefined set of colors.

Here's a breakdown of the script:

Importing Libraries:
    The script imports necessary libraries, including logging for logging messages, MoviePy for video editing, PIL (Pillow) for image processing, NumPy for numerical operations, scikit-learn for k-means clustering, and the os and random modules.

Configuring Logging:
    Logging is configured to record messages at the DEBUG level in a file named 'script_log.log'.

python

logging.basicConfig(filename='script_log.log', level=logging.DEBUG)

Defining the apply_kmeans Function:

This function takes an image and a set of colors and applies k-means clustering to segment the image into distinct color clusters.

Defining the Color Generator Function (ColA):

ColA generates a random RGB color.

python

def ColA():
Rr = random.randint(0, 255)
Gg = random.randint(0, 255)
Bb = random.randint(0, 255)
return [Rr, Gg, Bb]

Defining the Set of Colors:

The colors list is defined, and it contains several randomly generated RGB colors.

python

colors = [
ColA(),
ColA(),
ColA(),
ColA(),
ColA(),
ColA(),
ColA(),
ColA(),
ColA()
]

Setting Input and Output Paths:

Paths for the input video and the output segmented video are specified.

python

input_video_path = "/home/jack/Desktop/vids/6bd83e3c-2f9f-49bc-ac9f-ee233870b241.mp4"
output_video_path = "junk/video_segmented.mp4"

Creating a Temporary Directory:

A temporary directory is created to store intermediate files.

python

temp_directory = "temp_directory"
os.makedirs(temp_directory, exist_ok=True)

Processing Each Frame of the Video:

The script reads each frame of the input video, applies k-means clustering, and stores the segmented frames in a list.

python

clip = VideoFileClip(input_video_path)
segmented_frames = []

kmeans = KMeans(n_clusters=len(colors), random_state=42)
for i, frame in enumerate(clip.iter_frames(fps=clip.fps)):
if i == 0:
kmeans.fit(np.array(frame).reshape((-1, 3)))

labels = kmeans.predict(np.array(frame).reshape((-1, 3)))
segmented_frame_array = np.array(colors)[labels].reshape(np.array(frame).shape)
segmented_frame = Image.fromarray(np.uint8(segmented_frame_array))
segmented_frames.append(np.array(segmented_frame))

Defining a Function to Generate Frames:

A function (make_frame) is defined to generate frames for the segmented video.

python

def make_frame(t):
return segmented_frames[int(t * clip.fps)]

Creating the Segmented Video Clip:

A video clip (segmented_clip) is created using the VideoClip class, and the make_frame function is used to generate frames.

python

segmented_clip = VideoClip(make_frame, duration=clip.duration)

Writing the Segmented Video to a File:

The segmented video is written to an output file using the specified codec and audio codec.

python

segmented_clip.write_videofile(output_video_path, codec="libx264", audio_codec="aac", fps=clip.fps)

Cleaning Up Temporary Files:

The temporary directory is cleaned up by removing its contents.

python

for file_path in os.listdir(temp_directory):
file_path = os.path.join(temp_directory, file_path)
os.remove(file_path)
os.rmdir(temp_directory)

Logging Completion Message:

A completion message is logged.

python

logging.info("Script execution completed.")

In summary, the script takes a video, applies k-means clustering to each frame, assigns colors from a predefined set, and generates a segmented video. The process is logged, and temporary files are cleaned up.