PseudoRayTracingOSM/dataset_factory.py at master · djankov/PseudoRayTracingOSM · GitHub

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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

from sklearn.model_selection import train_test_split
import numpy as np
import torch
from torch.utils.data import Dataset
from torchvision import transforms
import h5py
from skimage import io, transform
import os
import matplotlib.pyplot as plt
from utils.invert import Invert
from pathloss_38901 import pathloss_38901

def dataset_factory(use_images=True, image_folder="images/snap_dk_250_png", transform=True, data_augment_angle=10):
    #Longitude,Latitude,Speed,Distance,Distance_x,Distance_y,PCI_64,PCI_65,PCI_302
    selected_features = [0, 1, 3, 4, 5, 6, 7, 8] #
     # ['SINR', 'RSRP', 'RSRQ', 'Power']
    selected_targets = [1]
    dataset_path='dataset'
    features = np.load("{}\\training_features.npy".format(dataset_path))
    targets = np.load("{}\\training_targets.npy".format(dataset_path))
    test_features =  np.load("{}\\test_features.npy".format(dataset_path))
    test_targets = np.load("{}\\test_targets.npy".format(dataset_path))
    target_mu = np.load("{}\\targets_mu.npy".format(dataset_path))
    target_std = np.load("{}\\targets_std.npy".format(dataset_path))
    features_mu = np.load("{}\\features_mu.npy".format(dataset_path))
    features_std = np.load("{}\\features_std.npy".format(dataset_path))
    images = np.load("{}\\train_image_idx.npy".format(dataset_path))
    test_images = np.load("{}\\test_image_idx.npy".format(dataset_path))


    features = features[:, selected_features]
    test_features = test_features[:, selected_features]
    features_mu = features_mu[selected_features]
    features_std = features_std[selected_features]


    targets = targets[:, selected_targets]
    test_targets = test_targets[:, selected_targets]
    target_mu = target_mu[selected_targets]
    target_std = target_std[selected_targets]


    # Data augmentation
    if transform:
        #composed = transforms.Compose([transforms.ToPILImage(), transforms.Grayscale(), transforms.RandomAffine(data_augment_angle, shear=10), transforms.ToTensor()])
        composed = transforms.Compose([transforms.ToPILImage(),  transforms.Grayscale(), Invert(), transforms.RandomAffine(data_augment_angle, shear=10), transforms.ToTensor()])
    else:
        composed = None

    # Dataset
    train_dataset = DrivetestDataset(features, targets, images, target_mu, target_std, features_mu, features_std, use_images, image_folder, transform=composed)
    #valid_dataset = DrivetestDataset(images, features, targets, valid_idx, target_mu, target_std, features_mean, features_std, use_images, image_folder)
    test_dataset = DrivetestDataset(test_features, test_targets, test_images,  target_mu, target_std, features_mu, features_std, use_images, image_folder, transform=transforms.Compose([transforms.ToPILImage(), transforms.Grayscale(), transforms.ToTensor()]))
    return train_dataset, test_dataset


class DrivetestDataset(Dataset):
    def __init__(self, features, targets, images, target_mu, target_std, feature_mu, feature_std, use_images, image_folder, transform=None):
        self.features = features
        self.targets = targets
        self.image_idx = images
        self.target_mu = target_mu
        self.target_std = target_std
        self.feature_mu = feature_mu
        self.feature_std = feature_std
        self.distances = (self.features[:,2] * self.feature_std[2])+self.feature_mu[2]
        self.targets_unnorm = (self.targets * self.target_std)+self.target_mu
        self.use_images = use_images
        self.image_folder = image_folder
        self.transform = transform
        self.image_size = io.imread(os.path.join(self.image_folder, "{}.png".format(0))).shape

    def get_811Mhz_idx(self):
        return np.argwhere(np.asarray(self.features[:,7] != 1))

    def get_2630Mhz_idx(self):
        return np.argwhere(np.asarray(self.features[:,7] == 1))

    def __getitem__(self, index):
        idx = self.image_idx[index]
        X = torch.from_numpy(self.features[index]).float() # Features (normalized)
        if self.use_images:
            if self.image_folder == None: #images are then pointer to hdf5
                image = self.image_idx[index]
            else:
                img_name = os.path.join(self.image_folder, "{}.png".format(idx))
                image = io.imread(img_name)
                image = image / 255
            A = torch.from_numpy(image).float().permute(2,0,1)

        else:
            A = torch.tensor(0)
        y = torch.from_numpy(self.targets[index]).float() # Target
        dist = torch.abs(torch.tensor(self.distances[index])).float().view(1) # Unormalized distance
        dist = dist * 1000 # to meters


        if self.use_images:
            if self.transform:
                A = self.transform(A)

        return X, A, y, dist

    def __len__(self):
        return len(self.features)

if __name__ == '__main__':
    train, test = dataset_factory()
    data = train.__getitem__(1)

    fig = plt.figure(figsize=(5,5))
    plt.imshow(data[1].permute(1,2,0).numpy())
    plt.show()