Video & Image Colorization Using OpenCV Python | AI COLORIZATION FOR IMAGES & VIDEOS

Introduction

In this article, we will explore how to colorize old black and white images and videos using OpenCV with Python. Additionally, we will implement CUDA to speed up the colorization process for videos. Let’s get started!

Setting Up the Environment

Before we dive into the code, make sure you have the necessary files. You will need an empty directory containing two Python files: colorizer.py, which will implement our class, and main.py, which will act as our driver program. Also, download the pre-trained model shared by Rich Yang—details and a link can be found in the description. Once you have everything set up, ensure your project structure includes an input folder for black and white media and an output folder for saving the colorized results.

The Colorizer Class

Initializing the Class

1. Imports: Start by importing essential packages including NumPy, OpenCV (cv2), and the os module functions.
2. Class Declaration: Create a class named Colorizer with an __init__ function that accepts height and width for default resolutions.
3. Load the Model: Use OpenCV's DNN module to read the pre-trained model. As the model is a Caffe model, we need readNetFromCaffe to load both the prototxt and the binary model files.

Preprocessing the Model

• Load the pre-trained cluster centers from a NumPy dump and set them to a specific layer.
• Adjust another layer with specific values to prepare the model for input data.

Image Processing

Colorization Method

1. Define a method called process_image, which will accept an image path.

2. Inside this method:

   • Read and store the image.
   • Resize the image according to the provided dimensions.

3. Create another method named process_frame to manage the colorization:

   • Swap the BGR to RGB channels since OpenCV uses BGR format while the model expects RGB.
   • Normalize the image and convert the color space from RGB to Lab, extracting the L channel which contains luminance information.
   • Prepare the L channel for input into the model by resizing it to the expected 224x224 size and adjusting its values.

4. Feed the processed L channel into the model and perform a forward pass to get a 256x256 output. Resize this output back to the original dimensions, concatenate the L channel with the colorized output, and perform denormalization to finalize the colored image.

Saving and Displaying Results

• Combine the original black and white image with the colorized version for a side-by-side comparison. Save the final result into the output directory.

Main Program Execution

In the main.py file, import the Colorizer class and create an instance. Call the process_image method with any black and white image from the input folder and run the program to view the results.

Video Colorization

Implementing Video Processing

1. Create a method for processing video files. This function will:

   • Attempt to read the video, and return an error message if it fails.
   • Read the first frame and set up the video writer, ensuring it maintains the original frame rate.

2. Loop through the frames, resizing and processing each one to create a colorized output. Display frames with the calculated FPS on-screen.

Optimizing with CUDA

If you are using OpenCV built with CUDA support, you can achieve faster processing speeds by setting two specific backend parameters within the class. This can significantly increase the FPS from approximately 2-3 to between 12-13.

Conclusion

After implementing the above steps, you can colorize both images and videos with relative ease. Try experimenting with your own media files, and let us know your experiences in the comments. Don't forget to like this article, and we hope to see you again in our future projects!

Keywords

• Image Colorization
• OpenCV
• Python
• CUDA
• Black and White
• DNN
• Pre-trained Model
• Video Processing

FAQ

Q: What is the purpose of the Colorizer class?
A: The Colorizer class is designed to handle the processes of image and video colorization using a pre-trained neural network model in OpenCV.

Q: Can I use the code without CUDA?
A: Yes, CUDA is optional. The code will work without it, but you may experience slower processing speeds for video colorization.

Q: Where can I find the pre-trained model mentioned in the article?
A: The model is shared by Rich Yang, and a link to it can be found in the description of the related video or the corresponding project repository.

Q: What file formats can I use for colorization?
A: You can use common formats such as JPEG, PNG for images and MP4, AVI for videos as inputs for the colorization process.

Q: How do I save the colorized results?
A: The results are automatically saved in the output directory defined in your project structure after processing the images or videos.