Với sự phát triển của AI trong những năm gần đây, càng ngày càng có nhiều các phần mềm chỉnh sửa hỉnh ảnh tốt hơn, giúp chúng ta tiết kiệm thời gian hơn so với sử dụng các công cụ chuyên dụng như Photoshop. Một số ứng dụng hay của AI trong mảng này là thêm màu cho ảnh đen trắng (Image Colorization), thêm khung hình cho video (Video Frame Interpolation), tăng độ phân giải của video (Super Resolution), …

Gần đây, mình đã thử nghiệm mô hình tự động thêm màu cho ảnh đen trắng từ bài báo Colorful Image Colorization(Image Colorization) và rất ấn tượng với những bức ảnh màu đạt được. Do đó, trong bài hôm nay, hãy cùng thử tìm hiểu về mô hình này và cách áp dụng nó cho những bức ảnh của chúng ta.

Colorful Image Colorization

Bài toán thêm màu cho ảnh đen trắng không phải là mới trong thị giác máy tính, nhưng phải đến năm 2016 khi Colorful Image Colorization được công bố, chúng ta mới biết đến một mô hình hoạt động tốt mà không cần đến sự điều chỉnh từ con người. Dưới đây là một vài ví dụ lấy từ bài báo gốc, cho thấy khả năng “tô màu” của mô hình này (mình mà tô chắc chắn sẽ xấu hơn).

Ví dụ về ảnh đen trắng đầu vào và kết quả (ảnh từ bài báo gốc)

Colorful Image Colorization hoạt động như thế nào

Dự đoán màu trong không gian màu CIELAB

Đầu tiên, các tác giả dùng ảnh biểu diễn trong không gian màu CIELAB thay vì RGB. Không gian màu CIELAB gồm 3 kênh: L, a & b; trong đó kênh L chứa thông tin về độ đậm nhạt, còn kênh a và b chứa thông tin về màu sắc của ảnh. Dựa trên tính chất này, bất cứ bức ảnh màu nào cũng có thể được dùng để huấn luyện mô hình với đầu vào là ảnh đen trắng (kênh L), đầu ra là thông tin về màu sắc (kênh ab) của bức ảnh đó.

Khái quát về mô hình trong Colorful Image Colorization (ảnh từ bài báo gốc)

Đưa về bài toàn phân loại

Bên cạnh đó, để đơn giản hoá quá trình tính toán, thay vì dự đoán dải màu ab với giá trị liên tục, các tác giả chia thành 313 nhóm màu. Với cách chia này, nhiệm vụ dự đoán màu bây giờ được chuyển thành một bài toán phân loại đa thức, trong đó cứ mỗi điểm ảnh có 313 lớp màu (classes) để lựa chọn.

Chia thông tin màu thành các nhóm (quantization)

Dùng hàm mất mát đặc biệt

Một điểm cần lưu ý nữa đó là cách chọn hàm mất mát (loss function) để có kết quả chính xác hơn với sự phân bố không đồng đều của màu sắc trong tập dữ liệu được dùng.

Chọn hàm mất mát phù hợp với sự phân bố không đều của màu sắc

Ngoài những điểm trên đây, bài báo còn có một số chi tiết khác khác về mặt triển khai và mở rộng, nếu có hứng thú thì các bạn tìm đọc thêm nhé, còn giờ bắt tay vào thực hành thôi :)

Thực hành

Bài thực hành hôm nay sẽ dùng mô hình huấn luyện sẵn từ Coloful Image Colorization để thêm màu cho một ảnh đen trắng bất kì.

Cấu trúc code

Mã nguồn được dùng trong bài hôm nay được lưu tại Panda ML Github để tiện cho việc tham khảo. Cấu trúc code cụ thể như sau:

.
├── images
│   ├── butterfly.jpeg
│   ├── landscape.jpeg
│   └── portrait.jpeg
├── models
│   ├── colorization_deploy_v2.prototxt
│   ├── colorization_release_v2.caffemodel
│   └── pts_in_hull.npy
├── colorize_image.py
└── utils.py
  • File colorize_image.py là file code chính được sử dụng
  • File utils.py chứa các hàm được dùng trong colorize_image.py
  • Thư mục models chứa các files mô hình đã được huấn luyện
  • Thư mục images chứa các ảnh được dùng để thử nghiệm mô hình

Để áp dụng tô màu cho 1 ảnh mới, chúng ta có thể dùng đoạn code dưới đây:

python colorize_image.py --image [đường dẫn đến ảnh]

Nội dung code

Trước tiên chúng ta cần đọc file mô hình:

1
2
3
4
5
6
7
8
9
10
11
12
# load model
weights_path = "models/colorization_release_v2.caffemodel"
config_path = "models/colorization_deploy_v2.prototxt"
net = cv2.dnn.readNetFromCaffe(config_path, weights_path)

# add quantized_ab center, will be used for rebalancing
pts = np.load("models/pts_in_hull.npy")
pts = pts.transpose().reshape(2, 313, 1, 1).astype("float32")
class8 = net.getLayerId("class8_ab")
conv8 = net.getLayerId("conv8_313_rh")
net.getLayer(class8).blobs = [pts]
net.getLayer(conv8).blobs = [np.full([1, 313], 2.606, dtype="float32")]
  • Dòng 2 - 4: Đọc mô hình được huấn luyện sẵn (mô hình được huấn luyện trên nền tảng Caffe).
  • Dòng 7 - 12: Đọc thông tin về các lớp màu đã có sẵn, được dùng để cân bằng lại phân bố màu không đều trong mô hình.

Bước tiếp theo là load ảnh và thực hiện các bước tiền xử lý cho phù hợp với yêu cầu của mô hình:

1
2
3
4
5
6
7
8
9
10
# load image
image = cv2.imread(image_path)
height, width = image.shape[:2]
scaled = image.astype("float32") / 255.0
lab = cv2.cvtColor(scaled, cv2.COLOR_BGR2LAB)

# resize, extract L channel, perform mean centering
resized = cv2.resize(lab, (224, 224))
L = l_channel_from_lab_image(resized)
L -= 50
  • Dòng 2 - 5: Load ảnh và chuyển ảnh sang không gian màu CIELAB.
  • Dòng 8 - 10: Trích xuất kênh đậm nhạt (L channel).
    • Hàm l_channel_from_lab_image được dùng để trích xuất kênh L từ ảnh CIELAB (trong file utils.py)

Sau đó là áp dụng mô hình lên ảnh đen trắng đã được xử lý ở bước trên:

1
2
3
4
5
6
7
8
# run model
net.setInput(cv2.dnn.blobFromImage(L))
ab = net.forward()[0, :, :, :].transpose((1, 2, 0))

# construct colored image from original L channel & predicted ab channel
predicted_ab = cv2.resize(ab, (width, height))  # resize to original image size
original_L = l_channel_from_lab_image(lab)
colorized = rgb_from_l_and_ab(original_L, predicted_ab)
  • Dòng 2 - 3: Chạy mô hình và lấy kết quả kênh màu ab.
  • Dòng 6 - 8: Kết hợp độ đậm nhạt (original_L) và màu sắc dự đoán (predicted_ab) để có được ảnh màu.

Và cuối cùng là hiển thị ảnh đen trắng đầu vào cùng ảnh màu đã được tái hiện lại:

1
2
3
4
# display image
cv2.imshow("Original", image)
cv2.imshow("Colorized", colorized)
cv2.waitKey(0)

Mã nguồn

Panda ML Github

Kết quả

Dưới đây là kết quả khi áp dụng mô hình của chúng ta lên ảnh trong thư mực images với 3 chủ đề: ảnh động vật, phong cảnh & chân dung (nguồn ảnh: Bình Minh digital, websosanh.vn, báo dantri.com.vn):

Tạm kết

Mặc dù đã biết đến mô hình này từ lâu nhưng đến giờ mình mới có dịp thử nghiệm với nó và phải nói là kết quả còn trên cả mong đợi. Bước tiếp theo mình sẽ thử áp dụng vào video đen trắng để cho thêm phần sinh động, giống như AI phục chế màu cho video trăm tuổi.

Còn các bạn, các bạn thấy kết quả ra sao? Và các bạn thấy mô hình này còn có thể áp dụng như thế nào thì cho mình biết với nhé.

Tham khảo

Leave a comment