将gid数据集的标签转换成语义分割模式

Gid的标签是三维tif格式,需要将其转换成适合语义分割的模式。
代码如下:

import os

import cv2
import tifffile as tiff
import numpy as np
from PIL import Image

def RGB_to_labe15l(img):
    label = np.zeros(shape=[img.shape[0], img.shape[1]], dtype=np.uint8)
    for i in range(0, img.shape[0]):
        print(i)
        for j in range(0, img.shape[1]):
            if all(np.array(img[i][j]) == np.array([200, 0, 0])):  # industrial land 1
                label[i][j] = 1
            elif all(np.array(img[i][j]) == np.array([250, 0, 150])):  # urban residential  2
                label[i][j] = 2
            elif all(np.array(img[i][j]) == np.array([200, 150, 150])):  # rural residential 3
                label[i][j] = 3
            elif all(np.array(img[i][j]) == np.array([250, 150, 150])):  # traffic land 4
                label[i][j] = 4
            elif all(np.array(img[i][j]) == np.array([0, 200, 0])):  # paddy field 5
                label[i][j] = 5
            elif all(np.array(img[i][j]) == np.array([150, 250, 0])):  # irrigated land
                label[i][j] = 6
            elif all(np.array(img[i][j]) == np.array([150, 200, 150])):  # dry cropland
                label[i][j] = 7
            elif all(np.array(img[i][j]) == np.array([200, 0, 200])):  # garden plot
                label[i][j] = 8
            elif all(np.array(img[i][j]) == np.array([150, 0, 250])):  # arbor woodland
                label[i][j] = 9
            elif all(np.array(img[i][j]) == np.array([150, 150, 250])):  # shrub land
                label[i][j] = 10
            elif all(np.array(img[i][j]) == np.array([250, 200, 0])):  # natural grassland
                label[i][j] = 11
            elif all(np.array(img[i][j]) == np.array([200, 200, 0])):  # artificial grassland
                label[i][j] = 12
            elif all(np.array(img[i][j]) == np.array([0, 0, 200])):  # river
                label[i][j] = 13
            elif all(np.array(img[i][j]) == np.array([0, 150, 200])):  # lake
                label[i][j] = 14
            elif all(np.array(img[i][j]) == np.array([0, 200, 250])):  # pound
                label[i][j] = 15
            else:
                label[i][j] = 0
            pass
    return label

def RGB_to_label5(img):
    label = np.zeros(shape=[img.shape[0], img.shape[1]], dtype=np.uint8)
    for i in range(0, img.shape[0]):
        print(i)
        for j in range(0, img.shape[1]):
            if all(np.array(img[i][j]) == np.array([255, 0, 0])):  # 红⾊,建成区
                label[i][j] = 1
            elif all(np.array(img[i][j]) == np.array([0, 255, 0])):  # 绿⾊,农⽤地
                label[i][j] = 2
            elif all(np.array(img[i][j]) == np.array([0, 255, 255])):  # 天蓝⾊,林地
                label[i][j] = 3
            elif all(np.array(img[i][j]) == np.array([255, 255, 0])):  # 黄⾊,草地
                label[i][j] = 4
            elif all(np.array(img[i][j]) == np.array([0, 0, 255])):  # 蓝⾊,⽔系
                label[i][j] = 5
            else:
                label[i][j] = 0
            pass
    return label
#   批量转换 将三维标签转换为一维
path = './data/gid/label_5classes/'
savedpath = './data/gid/annotations/labels/'
filelist = os.listdir(path)   # 带后缀名的list
print(filelist)

for item in filelist:
    print("开始转换:", item)
    im = tiff.imread(path + item)  # 打开图片
    label = RGB_to_label5(im)
     print(label.shape)
     save_path = os.path.join(savedpath, os.path.splitext(item)[0]+ '.png')
     label = Image.fromarray(np.uint8(label))
     label.save(save_path)


你可能感兴趣的:(计算机视觉,python,opencv)