Python处理图片并实现生成天际线

#-*- coding:utf-8 -*-
import sys
from os.path import exists
import cv2
import numpy as np

def getImage(height, width, channels):
    image = np.zeros([height, width, 3], np.uint8) # 三通道顺序是BGR
    # 三层循环逐个修改像素点
    for row in range(height):
        for col in range(width):
                for c in range(channels):
                    image[row, col, c] = 0
    return image

def isWhite(pixel_value, threshold): #阈值可以取10、20、30、50、100
    res = False
    if pixel_value[0] > threshold and pixel_value[1] > threshold and pixel_value[2] > threshold: # 10、10、10 50、50、50 这里是天空和地面楼山的分界线，需要调参
        res = True
    return res

def isPureWhite(pixel_value):
    res = False
    if pixel_value[0] == 255 and pixel_value[1] == 255 and pixel_value[2] == 255: # >3|>3|>3 10、10、10
        res = True
    return res

def getRowNumberSpecificCol(image, col):
    res_row = -1
    height, width = image.shape[0:2]
    if col >= 0 and col < width:
        for row in range(0, height):
            pv = image[row][col]
            if(pv[0] > 0 and pv[1] > 0 and pv[2] >0):
                res_row = row
                break
    return res_row

def getEnhancedEdgeImageFromEdgeImage(edge_Image):
    edge_SrcImage = edge_Image
    height, width = edge_SrcImage.shape[0:2]
    for col in range(1, width):
        for row in range(0, height):
            pixel_value = edge_SrcImage[row][col]  # 计算红绿蓝三波段的平均值
            if isPureWhite(pixel_value):
                r_last = getRowNumberSpecificCol(edge_SrcImage, col - 1)
                if r_last:
                    if row > r_last:
                        minR, maxR = r_last, row
                        for k in range(minR, maxR):
                            edge_SrcImage[k][col - 1][0] = 255
                            edge_SrcImage[k][col - 1][1] = 255
                            edge_SrcImage[k][col - 1][2] = 255
                    else:
                        minR, maxR = row, r_last
                        for k in range(minR, maxR):
                            edge_SrcImage[k][col][0] = 255
                            edge_SrcImage[k][col][1] = 255
                            edge_SrcImage[k][col][2] = 255
    # cv2.imshow("Enhanced-edge-image", edge_SrcImage)
    return edge_SrcImage

def getFileExtensionname(filename):
    res = ".png"
    dot_index = -1
    for i in range(len(filename), 0):
        if filename[i] == '.':
            dot_index = i
            break
    if dot_index != -1:
        res = filename[dot_index: len(filename)-1]
    return res

if __name__ == '__main__':
    origin_pic_filename = "D:/test.png"
    sky_ground_threshold = 30
    isDownSampling = False
    if (len(sys.argv) == 1):
        print(sys.argv[0])
        origin_pic_filename = ""
    elif(len(sys.argv) == 2):
        origin_pic_filename = str(sys.argv[1])
    elif(len(sys.argv) == 3):
        origin_pic_filename = str(sys.argv[1])
        sky_ground_threshold = int(sys.argv[2])
    elif (len(sys.argv) == 4):
        origin_pic_filename = str(sys.argv[1])
        sky_ground_threshold = int(sys.argv[2])
        if(int(sys.argv[3]) == 1):
            isDownSampling = True
    if origin_pic_filename != "" and sky_ground_threshold > 0:
        print(("输入图片文件名为：{0}").format(origin_pic_filename))
        print(("天空地面分界灰度阈值为：{0}").format(sky_ground_threshold))
        suffix_name = getFileExtensionname(origin_pic_filename)
        print(("后缀名为：{0}").format(suffix_name))

        srcImage = cv2.imread(origin_pic_filename)
        inputSrcImage = srcImage
        if isDownSampling:
            inputSrcImage = cv2.pyrDown(inputSrcImage)
        height, width = inputSrcImage.shape[0:2]
        print(("高度：{0}, 宽度：{1}").format(height, width))
        cv2.namedWindow('downsampling-image', cv2.WINDOW_AUTOSIZE)
        cv2.imshow("downsampling-image", inputSrcImage)
        Sobelx = cv2.Sobel(inputSrcImage, cv2.CV_64F, 1, 0)
        Sobely = cv2.Sobel(inputSrcImage, cv2.CV_64F, 0, 1)
        Sobelx = cv2.convertScaleAbs(Sobelx)
        Sobely = cv2.convertScaleAbs(Sobely)
        # cv2.imshow("sobel-x-Abs", Sobelx)
        # cv2.imshow("sobel-y-Abs", Sobely)
        Sobelxy = cv2.addWeighted(Sobelx, 0.5, Sobely, 0.5, 0)
        cv2.namedWindow('sobel-xy', cv2.WINDOW_AUTOSIZE)
        cv2.imshow('sobel-xy', Sobelxy)
        edgeImage = getImage(height, width, 3)
        for col in range(0, width):
            for row in range(0, height):
                pixel_value = Sobelxy[row][col]  # 计算红绿蓝三波段的平均值
                if isWhite(pixel_value, sky_ground_threshold):
                    edgeImage[row][col][0] = 255
                    edgeImage[row][col][1] = 255
                    edgeImage[row][col][2] = 255
                    break
        cv2.namedWindow('edge-image', cv2.WINDOW_AUTOSIZE)
        cv2.imshow('edge-image', edgeImage)
        cv2.imwrite(origin_pic_filename.replace(suffix_name, "-ZGetEdge.png"), edgeImage)
        enhanced_edgeImage = getEnhancedEdgeImageFromEdgeImage(edgeImage)
        cv2.namedWindow('enhanced-edge-image', cv2.WINDOW_AUTOSIZE)
        cv2.imshow('enhanced-edge-image', enhanced_edgeImage)
        cv2.imwrite(origin_pic_filename.replace(suffix_name, "-EnhancedEdge.png"), enhanced_edgeImage)

        for col in range(0, width):
            for row in range(0, height):
                pixel_value = enhanced_edgeImage[row][col]  # 计算红绿蓝三波段的平均值
                if isPureWhite(pixel_value):
                    if row+2 < height:
                        inputSrcImage[row+2][col][0] = 0
                        inputSrcImage[row+2][col][1] = 0
                        inputSrcImage[row+2][col][2] = 255
                    else:
                        inputSrcImage[row][col][0] = 0
                        inputSrcImage[row][col][1] = 0
                        inputSrcImage[row][col][2] = 255
                    # inputSrcImage[row][col][0] = 0
                    # inputSrcImage[row][col][1] = 0
                    # inputSrcImage[row][col][2] = 255
                    # break #最开始从每列遍历从上到下找第一个分界点就停止才用break

        cv2.namedWindow('RedEdge-image', cv2.WINDOW_AUTOSIZE)
        cv2.imshow('RedEdge-image', inputSrcImage)
        cv2.imwrite(origin_pic_filename.replace(suffix_name, "-RedEdge.png"), inputSrcImage)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
        print('Success!')
        cv2.waitKey()
        cv2.destroyAllWindows()

python GetSkyLine.py test.jpg  100

python GetSkyLine.py test.jpg  50  1