-
Notifications
You must be signed in to change notification settings - Fork 3
/
dataset.py
54 lines (41 loc) · 1.81 KB
/
dataset.py
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
# Dataset construction based on WHU-OHS
import os
import torch
from torch.utils import data
import numpy as np
from osgeo import gdal
class WHU_OHS_Dataset(data.Dataset):
def __init__(self, image_file_list, label_file_list, use_3D_input=False, channel_last=False):
self.image_file_list = image_file_list
self.label_file_list = label_file_list
self.use_3D_input = use_3D_input
self.channel_last = channel_last
# Statistics of samples of each class in the dataset
def sample_stat(self):
sample_per_class = torch.zeros([24])
for label_file in self.label_file_list:
label = gdal.Open(label_file, gdal.GA_ReadOnly)
label = label.ReadAsArray()
count = np.bincount(label.ravel(), minlength=25)
count = count[1:25]
count = torch.tensor(count)
sample_per_class = sample_per_class + count
return sample_per_class
def __len__(self):
return len(self.image_file_list)
def __getitem__(self, index):
image_file = self.image_file_list[index]
label_file = self.label_file_list[index]
name = os.path.basename(image_file)
image_dataset = gdal.Open(image_file, gdal.GA_ReadOnly)
label_dataset = gdal.Open(label_file, gdal.GA_ReadOnly)
image = image_dataset.ReadAsArray()
label = label_dataset.ReadAsArray()
if(self.channel_last):
image = image.transpose(1, 2, 0)
# The image patches were normalized and scaled by 10000 to reduce storage cost
image = torch.tensor(image, dtype=torch.float) / 10000.0
if(self.use_3D_input):
image = image.unsqueeze(0)
label = torch.tensor(label, dtype=torch.long) - 1
return image, label, name