将黑色mask图片转为彩色标注文件

从黑色掩码图片可视化为彩色标注图像        

1. 读取文件夹中的图片信息，生成csv文档
2. 从csv文档中一次读取一个对象
3. 批量生成可视化文件 

1、读取文件夹中的图片信息，生成csv文档

 只需修改地址信息及存入文档信息

image_mask.py

import os
import csv
import re

# 要读取的文件的根目录
# root_path = 'LoveDA/Train/Urban/masks_png' #待读取的训练集地址
root_path = 'G:/LoveDA/Val/Urban/masks_png'  #待读取的验证集地址

# 将所有目录下的文件信息放到列表中
def get_Write_file_infos(path):
    # 文件信息列表
    file_infos_list = []
    # 遍历并写入文件信息
    for root, dirnames, filenames in os.walk(path):

        for filename in filenames:
            file_infos = {}
            # dirname = root
            # dirname = re.findall(r'val\\(.+)', dirname)
            # 正则表达式来截取文件夹名称
            # 正则表达式还得学习
            filename1 = filename.split('.png')[0]
            file_infos["file"] = filename1
            # file_infos["species"] = dirname[0]
            # file_infos["图片"] = ''

            # 将数据追加字典到列表中
            file_infos_list.append(file_infos)

    return file_infos_list


# 写入csv文件
def write_csv(file_infos_list):
    # with open('data1.csv', 'a+', newline='') as csv_file:  #读取的训练集名称存入data1
    with open('data2.csv', 'a+', newline='') as csv_file:    ##读取的训练集名称存入data2
        # csv_writer = csv.DictWriter(csv_file, fieldnames=['file', 'species'])
        csv_writer = csv.DictWriter(csv_file, fieldnames=['file'])
        csv_writer.writeheader()
        for each in file_infos_list:
            csv_writer.writerow(each)


# 主函数
def main():
    # 调用获取文件信息的函数
    file_infos_list = get_Write_file_infos(root_path)
    # 执行写入程序
    write_csv(file_infos_list)


# 主程序入口

if __name__ == '__main__':
    main()

补充：生成list文档的代码如下：

import fnmatch
import os
import pandas as pd
import numpy as np
import sys


InputStra = 'LoveDA/Train/Urban/masks_png' #待读取文件的地址
InputStrb = '*.png'  #在每个名字后，加入png后缀

def ReadSaveAddr(Stra,Strb):
    #print(Stra)
    #print(Strb)
    print("Read :",Stra,Strb)
    a_list = fnmatch.filter(os.listdir(Stra),Strb)
    print("Find = ",len(a_list))
    df = pd.DataFrame(np.arange(len(a_list)).reshape((len(a_list),1)),columns=['Addr'])
    df.Addr = a_list
    #print(df.head())
    df.to_csv('data1',columns=['Addr'],index=False,header=False)
    print("Write To data1 !")

ReadSaveAddr(InputStra,InputStrb)

2、从csv文档中一次读取一个对象

csvread.py

import csv
import os
# 读取数据
import pandas as pd
# data = pd.read_csv("data1.csv",header=None)   
data = pd.read_csv("data1.csv")  #直接将csv表中的第一行当作表头  训练集
# print(data)

# ①获取所有列，并存入一个数组中
import numpy as np
data = np.array(data)
# print(data)
# [[1366]
#  [1367]
#  [1368]
#  ...
#  [2519]
#  [2520]
#  [2521]]  二维数组

for element in data:   #element=[3066]
    ele_str = str(element[0])  #首先读取二维数组，element是一维数组，读取里面的元素，然后转为字符串形式

3、生成可视化文件 

render.py

from PIL import Image
import numpy as np

COLOR_MAP = dict(
    IGNORE=(0, 0, 0),
    Background=(255, 255, 255),
    Building=(255, 0, 0),
    Road=(255, 255, 0),
    Water=(0, 0, 255),
    Barren=(159, 129, 183),
    Forest=(0, 255, 0),
    Agricultural=(255, 195, 128),
)


def render(mask_path, vis_path):
    new_mask = np.array(Image.open(mask_path)).astype(np.uint8)
    cm = np.array(list(COLOR_MAP.values())).astype(np.uint8)
    color_img = cm[new_mask]
    color_img = Image.fromarray(np.uint8(color_img))
    color_img.save(vis_path)


if __name__ == '__main__':
    mask_path = r'G:\LoveDA\Train\Urban\masks_png\1366.png'  #待读取的地址（str）
    vis_path = r'C:\Users\28123\Desktop\1366_vis.png'  #存入的文档
    render(mask_path, vis_path) 

4、批量生成！！！

csvread.py 

难点：从csv中读取对象，转为数组，然后再转为字符串，形成路径，再从路径中获取图片。 

import csv
import os
# 读取数据
import pandas as pd
data = pd.read_csv("data1.csv")  #直接将csv表中的第一行当作表头，有一个‘file’对象


# ①获取所有列，并存入一个数组中
import numpy as np
data = np.array(data)

from PIL import Image
import numpy as np
COLOR_MAP = dict(
        IGNORE=(0, 0, 0),
        Background=(255, 255, 255),
        Building=(255, 0, 0),
        Road=(255, 255, 0),
        Water=(0, 0, 255),
        Barren=(159, 129, 183),
        Forest=(0, 255, 0),
        Agricultural=(255, 195, 128),
    )


def render(mask_path, vis_path):
    new_mask = np.array(Image.open(mask_path)).astype(np.uint8)
    cm = np.array(list(COLOR_MAP.values())).astype(np.uint8)
    color_img = cm[new_mask]
    color_img = Image.fromarray(np.uint8(color_img))
    color_img.save(vis_path)



# 遍历数组
root = 'G:\\LoveDA\\Train\\Urban\masks_png\\'
root_vis = 'G:\\LoveDA\\Train\\Urban\\masks_vis\\'

for element in data:
    ele_str = str(element[0])
    if __name__ == '__main__':
        path_root = os.path.join(root, ele_str + '.png')
        vis_root = os.path.join(root_vis, ele_str + '.png')
        mask_path = path_root
        vis_path = vis_root
        render(mask_path, vis_path)