Python+OpenCV实现黑白老照片上色功能

import numpy as np
import cv2 as cv

def img_translate(img, tx, ty):
    """对图像进行平移"""
    heigh, width = img.shape[:2]
    m = np.float32([[1, 0, tx], [0,1, ty]])
    res = cv.warpAffine(img, m, (width, heigh))
    return res

def ssd(I1, I2):
    """ssd函数，衡量颜色通道是否对齐"""
    return np.sum((I1 - I2)*(I1 - I2))


def find_xy(img1, g):
    """找到最佳的平移参数，并对该颜色通道图片进行平移"""
    # 初始化
    loss = ssd(img1, g)
    img = img1
    u = 0
    v = 0
    # 根据ssd函数寻找最佳的对齐位置
    for i in range(-20, 30):
        for j in range(-20, 30):
            img2 = img_translate(img1, i, j)
            loss1 = ssd(img2, g)
            if loss > loss1:
                loss = loss1
                img = img2
                u = i
                v = j
    print(u, v)
    return img


def readImage(imname):
    """read in the image"""
    im = cv.imread(imname)
    im = cv.cvtColor(im, cv.COLOR_BGR2GRAY)
    return im

def separate(im):
    """separate color channels"""
    # compute the height of each part (just 1/3 of total)
    height = np.floor(im.shape[0] / 3.0).astype(np.int)
    b = im[:height]
    g = im[height: 2 * height]
    r = im[2 * height: 3 * height]
    return b, g, r

def merge(b,g,r):
    """将三个颜色通道进行merge"""
    return cv.merge((b,g,r))

def gaussianPyramid(img):
    """直接使用OpenCV的高斯金字塔进行实现"""
    return cv.pyrDown(img)

if __name__ == '__main__':
    # name of the input file
    #imname = 'images/train.tif'
    imname = 'images/three_generations.tif'
    #imname = 'images/lady.tif'
    #imname = 'images/emir.tif'
    #imname = 'images/icon.tif'
    #imname = 'images/self_portrait.tif'
    #imname = 'images/village.tif'
    #imname = 'images/turkmen.tif'

    im = readImage(imname)
    print(im.shape)
    #cv.imshow("source image", im)

    # 获取平均切割的三个颜色通道
    b, g, r = separate(im)

    # 对三个颜色通道分别应用高斯金字塔
    b = gaussianPyramid(b)
    b = gaussianPyramid(b)
    b = gaussianPyramid(b)

    g = gaussianPyramid(g)
    g = gaussianPyramid(g)
    g = gaussianPyramid(g)

    r = gaussianPyramid(r)
    r = gaussianPyramid(r)
    r = gaussianPyramid(r)

    im_out0 = merge(b, g, r)
    cv.imshow("before", im_out0)
    # 颜色通道平移进行对齐，对齐的过程中以绿色作为基准
    b = find_xy(b, g)

    r = find_xy(r, g)
    # 将平移处理后的三通道merge，得到处理后的图片im_out1
    im_out1 = merge(b, g, r)

    # 将修复后的图片写进磁盘
    #cv.imwrite('out/after_' + imname[7:], im_out1)
    cv.imshow("after", im_out1)
    #print(im_out1.shape)
    cv.waitKey(0)