无人机语义分割数据集（Semantic Drone Dataset）使用方法

数据集介绍

Semantic Drone Dataset数据集下载地址
该语义无人机数据集专注于城市场景的语义理解，以提高无人机自主飞行和着陆程序的安全性。该图像描绘了在距地面 5 至 30 米的高度从最低点（鸟瞰）视角拍摄的 20 多座房屋。高分辨率相机用于获取尺寸为 6000x4000px (24Mpx) 的图像。训练集包含 400 张公开可用的图像，测试集包含 200 张私有图像。

该语义分割数据集包括20个种类：

• 树、草、其他植被、污垢、碎石
• 岩石、水、铺砌面积、水池、人
• 狗、车、自行车、屋顶、墙
• 栅栏、栅栏杆、窗户、门、障碍

数据集样本示例

数据集使用

生成掩模图

由于该数据集未提供训练集使用的掩模图，因此需要自己根据RGB值去还原掩膜图。可以通过以下代码转换成掩膜图或者给掩模图上色。

'''
colorTransformer.py
'''
import numpy as np
class ColorTransformer:
  def __init__(self):
    # color table.
    self.clr_tab = self.createColorTable()
    print(self.clr_tab)
    print(self.clr_tab.keys())
    # id table.
    id_tab = {}
    for k, v in self.clr_tab.items():
        id_tab[k] = self.clr2id(v)
    self.id_tab = id_tab
    print(self.id_tab)

  def createColorTable(self):
    clr_tab = {}
    clr_tab['unlabeled'] = [0, 0, 0]
    clr_tab['paved-area'] = [128, 64, 128]
    clr_tab['dirt'] = [130, 76, 0]
    clr_tab['grass'] = [0, 102, 0]
    clr_tab['gravel'] = [112, 103, 87]
    clr_tab['water'] = [28, 42, 168]
    clr_tab['rocks'] = [48, 41, 30]
    clr_tab['pool'] = [0, 50, 89]
    clr_tab['vegetation'] = [107, 142, 35]
    clr_tab['roof'] = [70, 70, 70]
    clr_tab['wall'] = [102, 102, 156]
    clr_tab['window'] = [254, 228, 12]
    clr_tab['door'] = [254, 148, 12]
    clr_tab['fence'] = [190, 153, 153]
    clr_tab['fence-pole'] = [153, 153, 153]
    clr_tab['person'] = [255, 22, 96]
    clr_tab['dog'] = [102, 51, 0]
    clr_tab['car'] = [9, 143, 150]
    clr_tab['bicycle'] = [119, 11, 32]
    clr_tab['tree'] = [51, 51, 0]
    clr_tab['bald-tree'] = [190, 250, 190]
    clr_tab['ar-marker'] = [112, 150, 146]
    clr_tab['obstacle'] = [2, 135, 115]
    clr_tab['conflicting'] = [255, 0, 0]
    return clr_tab

  def colorTable(self):
    return self.clr_tab
   
  def clr2id(self, clr):
    return clr[0]+clr[1]*255+clr[2]*255*255

  #transform to uint8 integer label
  def transform(self,label, dtype=np.int32):
    height,width = label.shape[:2]
    # default value is index of clutter.
    newLabel = np.zeros((height, width), dtype=dtype)
    id_label = label.astype(np.uint64)
    id_label = id_label[:,:,0]+id_label[:,:,1]*255+id_label[:,:,2]*255*255
    for tid,key in enumerate(self.clr_tab.keys()):
      val = self.id_tab[key]
      mask = (id_label == val)
      newLabel[mask] = tid
    return newLabel

  #transform back to 3 channels uint8 label
  def inverse_transform(self, label):
    label_img = np.zeros(shape=(label.shape[0], label.shape[1],3),dtype=np.uint8)
    values = list(self.clr_tab.values())
    for tid,val in enumerate(values):
      mask = (label==tid)
      label_img[mask] = val
    return label_img

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73

'''
prepareTrainIdFiles.py
'''
import os
import os.path as osp
import numpy as np
from colorTransformer import ColorTransformer
from PIL import Image

clrEnc = ColorTransformer()
def prepareTrainIDForDir():

    lbl_paths = "training_set/gt/semantic/label_images/"
    saveDirPath = "training_set/gt/semantic/label_TrainId/"
    for lbl_p in os.listdir(lbl_paths):
        print(lbl_p)
        lbl_path = lbl_paths + lbl_p
        trainId_path = saveDirPath + lbl_p
        gt = np.array(Image.open(lbl_path))
        trainId = clrEnc.transform(gt, dtype=np.uint8)
        Image.fromarray(trainId).save(trainId_path)

if __name__=='__main__':
    prepareTrainIDForDir()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

需要根据自己存放的路径调整lbl_paths，并且需要提前建立好掩模图像的保存文件夹，根据路径调整saveDirPath变量。
生成的掩模图如下：

如果想从掩膜图片中恢复出label图，可以通过以下代码：

'''
transformertest.py
'''
import os
import os.path as osp
import numpy as np
from colorTransformer import ColorTransformer
from PIL import Image

clrEnc = ColorTransformer()
def prepareTrainIDForDir():

    lbl_paths = "598.png"
    saveDirPath = "training_set/gt/semantic/label_TrainId/598.png"
    gt = np.array(Image.open(saveDirPath))
    trainId = clrEnc.inverse_transform(gt)
    Image.fromarray(trainId).save(lbl_paths)

if __name__=='__main__':
    prepareTrainIDForDir()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

使用前需要修改路径，此代码只实现了单张图上色，可以根据需要改成批量处理的代码。如根据上述掩码图重新生成标签图，得到如下图像，与ground truth 没有差别。

数据集加载 dataloader写法（基于pytorch）

'''
dataset.py
'''
# camera-ready
import torch
import torch.utils.data
import numpy as np
import cv2
import os

class DatasetTrain(torch.utils.data.Dataset):
    def __init__(self, base_dir):

        self.base_dir = base_dir
        self.img_dir = base_dir + "images/"
        self.label_dir = base_dir + "gt/semantic/label_TrainId/"

        self.new_img_h = 512
        self.new_img_w = 1024

        self.examples = []
        train_img_dir_path = self.img_dir
        label_img__dir_path = self.label_dir
        file_names = os.listdir(train_img_dir_path)

        for file_name in file_names:
            img_path = train_img_dir_path + file_name
            label_img_path = label_img__dir_path + file_name.split(".jpg")[0] + ".png"
            example = {}
            example["img_path"] = img_path
            example["label_img_path"] = label_img_path
            self.examples.append(example)

        self.num_examples = len(self.examples)

    def __getitem__(self, index):
        example = self.examples[index]

        img_path = example["img_path"]
        # print(img_path)
        img = cv2.imread(img_path, -1) 
        img = cv2.resize(img, (self.new_img_w, self.new_img_h),
                         interpolation=cv2.INTER_NEAREST) 
        label_img_path = example["label_img_path"]
        # print(label_img_path)
        label_img = cv2.imread(label_img_path, cv2.IMREAD_GRAYSCALE) 
        label_img = cv2.resize(label_img, (self.new_img_w, self.new_img_h),
                               interpolation=cv2.INTER_NEAREST) 

        # normalize the img (with the mean and std for the pretrained ResNet):
        img = img/255.0
        img = img - np.array([0.485, 0.456, 0.406])
        img = img/np.array([0.229, 0.224, 0.225]) 
        img = np.transpose(img, (2, 0, 1)) 
        img = img.astype(np.float32)

        # convert numpy -> torch:
        img = torch.from_numpy(img) 
        label_img = torch.from_numpy(label_img) 

        return (img, label_img)

    def __len__(self):
        return self.num_examples

if __name__ == "__main__":
    base_dir = "training_set/"
    train_dataset = DatasetTrain(base_dir = base_dir)
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                            batch_size=32, shuffle=True,
                                            num_workers=1,drop_last=True)
    from torch.autograd import Variable
    for step, (imgs, label_imgs) in enumerate(train_loader):
        imgs = Variable(imgs).cuda() # (shape: (batch_size, 3, img_h, img_w))
        # print(imgs.shape)
        label_imgs = Variable(label_imgs.type(torch.LongTensor)).cuda() # (shape: (batch_size, img_h, img_w))
        # print(label_imgs.shape)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

需要根据自己的路径设置base_dir变量。