Convert black mask picture to color annotation file

Visualize from black mask image to color annotation image         

1. Read the picture information in the folder , Generate csv file
2. from csv Read one object at a time in the document
3. Batch generation of visualization files  

1、 Read the picture information in the folder , Generate csv file

  Just modify the address information and save the document information

image_mask.py

import os
import csv
import re

#  The root directory of the file to be read 
# root_path = 'LoveDA/Train/Urban/masks_png' # The address of the training set to be read 
root_path = 'G:/LoveDA/Val/Urban/masks_png'  # Address of the verification set to be read 

#  Put the file information under all directories in the list 
def get_Write_file_infos(path):
    #  File information list 
    file_infos_list = []
    #  Traverse and write file information 
    for root, dirnames, filenames in os.walk(path):

        for filename in filenames:
            file_infos = {}
            # dirname = root
            # dirname = re.findall(r'val\\(.+)', dirname)
            #  Regular expressions to intercept folder names 
            #  Regular expressions have to be learned 
            filename1 = filename.split('.png')[0]
            file_infos["file"] = filename1
            # file_infos["species"] = dirname[0]
            # file_infos[" picture "] = ''

            #  Append data to the dictionary to the list 
            file_infos_list.append(file_infos)

    return file_infos_list


#  write in csv file 
def write_csv(file_infos_list):
    # with open('data1.csv', 'a+', newline='') as csv_file:  # The read training set name is saved in data1
    with open('data2.csv', 'a+', newline='') as csv_file:    ## The read training set name is saved in 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)


#  The main function 
def main():
    #  Call the function to get the file information 
    file_infos_list = get_Write_file_infos(root_path)
    #  Execute the writer 
    write_csv(file_infos_list)


#  Main program entry 

if __name__ == '__main__':
    main()

Add ： Generate list The code of the document is as follows ：

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


InputStra = 'LoveDA/Train/Urban/masks_png' # Address of the file to be read 
InputStrb = '*.png'  # After each name , Join in png suffix 

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、 from csv Read one object at a time in the document

csvread.py

import csv
import os
#  Reading data 
import pandas as pd
# data = pd.read_csv("data1.csv",header=None)   
data = pd.read_csv("data1.csv")  # Direct will csv The first row in the table acts as the header    Training set 
# print(data)

# ① Get all columns , And store it in an array 
import numpy as np
data = np.array(data)
# print(data)
# [[1366]
#  [1367]
#  [1368]
#  ...
#  [2519]
#  [2520]
#  [2521]]   Two dimensional array 

for element in data:   #element=[3066]
    ele_str = str(element[0])  # First read the two-dimensional array ,element It's a one-dimensional array , Read the elements inside , Then it is converted to string form 

3、 Generate visualization files  

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'  # Address to be read （str）
    vis_path = r'C:\Users\28123\Desktop\1366_vis.png'  # Saved documents 
    render(mask_path, vis_path) 

4、 Batch build ！！！

csvread.py 

difficulty ： from csv Read object in , Converted to an array , Then it is converted to a string , Form a path , Then get the picture from the path . 

import csv
import os
#  Reading data 
import pandas as pd
data = pd.read_csv("data1.csv")  # Direct will csv The first row in the table acts as the header , There is one ‘file’ object 


# ① Get all columns , And store it in an array 
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)



#  Traversal array 
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)