Python cv2灰度图转伪彩图

import cv2
import os
import pdb

data_path = './output/density_maps_mcnn_shtechA_660/'
file_name = 'output_IMG_1.png'
out_path = './out-temp/'
im_dir = os.path.join(data_path,file_name)
img = cv2.imread(im_dir)
#cv2.imshow('image',img)
pdb.set_trace()
ht = img.shape[0]
wd = img.shape[1]
for i in range(0,ht):
    for j in range(0,wd):
        B = img[i][j][0]
        G = img[i][j][1]
        R = img[i][j][2]
        if B <= 51:
            img[i][j][0] = 255
        elif B <= 102:
            img[i][j][0] = 255 - (B-51)*5
        elif B <= 153:
            img[i][j][0] = 0
        else : img[i][j][0] = 0

        if G <= 51:
            img[i][j][1] = G*5
        elif G <= 102:
            img[i][j][1] = 255
        elif G <= 153:
            img[i][j][1] = 255
        elif G <= 204:
            img[i][j][1] = 255 - int(128.0*(G-153.0)/51.0 + 0.5)
        else : img[i][j][1] = 127 - int(127.0*(G-204.0)/51.0 + 0.5)

        if R <= 51:
            img[i][j][2] = 0
        elif R <= 102:
            img[i][j][2] = 0
        elif R <= 153:
            img[i][j][2] = (R-102)*5
        elif G <= 204:
            img[i][j][2] = 255
        else : img[i][j][2] = 255
cv2.imwrite(os.path.join(out_path,file_name),img)