-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathdataset.py
166 lines (127 loc) · 5.94 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
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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
'''
Author: Jalen-Zhong [email protected]
Date: 2023-04-11 15:02:59
LastEditors: Jalen-Zhong [email protected]
LastEditTime: 2023-04-22 15:57:47
FilePath: \local ability of CNN\dataset.py
Description:
Reference or Citation:
Copyright (c) 2023 by [email protected], All Rights Reserved.
'''
import numpy as np
from random import randint
from tqdm import tqdm
from matplotlib import pyplot as plt
import random
import torch.utils.data as data
import h5py
import torch
from PIL import Image
def get_square_label(coords, pixels, rectangle_size):
x, y = coords[0], coords[1]
pixels[int(x - rectangle_size): int(x + rectangle_size), int(y - rectangle_size): int(y + rectangle_size)] = 255
density_map = pixels / 255
return density_map
# def get_square_label(coords, pixels, rectangle_size):
# x, y = coords[0], coords[1]
# pixels[int(x): int(x + rectangle_size), int(y): int(y + rectangle_size)] = 255
# density_map = pixels / 255
# return density_map
def save_h5(path, images, labels):
print('saving', path)
with h5py.File(name=path, mode='w') as file:
file['images'] = images
file['labels'] = labels
def load_h5(path):
print('loading', path)
file = h5py.File(name=path, mode='r')
return file['images'][:], file['labels'][:]
class DataFromH5File(data.Dataset):
def __init__(self, filepath):
h5File = h5py.File(filepath, 'r')
self.x = h5File['images']
self.y = h5File['labels']
def __getitem__(self, idx):
data = torch.from_numpy(self.x[idx]).float()
label = torch.tensor(self.y[idx])
return data, label
def __len__(self):
assert self.x.shape[0] == self.y.shape[0], "Wrong data length"
return self.x.shape[0]
class DataFromFileFolder(data.Dataset):
def __init__(self, file_list, transform=None):
self.file_list = file_list
self.transform = transform
self.filelength = len(file_list)
def __getitem__(self, idx):
img_path = self.file_list[idx]
img = Image.open(img_path)
img_transformed = self.transform(img)
tip = img_path.split("/")[-1].split(".")[0]
iszero = 1 if tip == 'cat' else 0
isone = 1 if tip == 'dog' else 0
label = [iszero, isone]
label = torch.tensor(label).float()
return img_transformed, label
def __len__(self):
return self.filelength
def generator(examples, image_size, rectangle_size, left_up_coors, right_down_coors, middle_coors, fix_coors, extra_lu_coors, extra_rd_coors):
list_images = []
list_labels = []
for i in tqdm(range(examples)):
left_up_pixels = np.zeros((image_size, image_size), np.uint8)
right_down_pixels = np.zeros((image_size, image_size), np.uint8)
middle_pixels = np.zeros((image_size, image_size), np.uint8)
fix_pixels = np.zeros((image_size, image_size), np.uint8)
extra_lu_pixels = np.zeros((image_size, image_size), np.uint8)
extra_rd_pixels = np.zeros((image_size, image_size), np.uint8)
left_up_seed = randint(0,1)
right_down_seed = randint(0,1)
middle_seed = randint(0,1)
XOR = ( left_up_seed ^ right_down_seed )
temp = np.int64(1) - XOR
label = temp
# iszero = 1 if temp == 0 else 0
# isone = 1 if temp == 1 else 0
# label = [iszero, isone]
fix_pixels = get_square_label(fix_coors, fix_pixels, rectangle_size)
if left_up_seed == 1:
left_up_pixels = get_square_label(left_up_coors, left_up_pixels, rectangle_size)
if right_down_seed == 1:
right_down_pixels = get_square_label(right_down_coors, right_down_pixels, rectangle_size)
# if middle_seed == 1:
middle_pixels = get_square_label(middle_coors, middle_pixels, rectangle_size)
# extra_lu_pixels = get_square_label(extra_lu_coors, extra_lu_pixels, rectangle_size)
# extra_rd_pixels = get_square_label(extra_rd_coors, extra_rd_pixels, rectangle_size)
image = left_up_pixels + right_down_pixels + fix_pixels + middle_pixels
image = image.reshape(1, image_size, image_size) / image.max()
list_images.append(image)
list_labels.append(label)
return list_images, list_labels
if __name__ == "__main__":
random.seed(2023)
train_examples = 10000
test_examples = 2000
image_size = 179
patch_number = 5
patch_size = image_size / patch_number
rectangle_size = patch_size / 4
fix_coors = [image_size - patch_size / 2, patch_size / 2]
# left_up_coors = [patch_size / 2, patch_size / 2]
# right_down_coors = [image_size - patch_size/2, image_size - patch_size/2]
left_up_coors = [patch_size - patch_size / 2 , patch_size / 2 ]
right_down_coors = [image_size - patch_size/2, image_size - patch_size/2]
middle_coors = [image_size / 2, image_size / 2]
extra_lu = [image_size/ patch_number, image_size/ patch_number]
extra_rd = [image_size - image_size/ patch_number, image_size - image_size/ patch_number]
# fix_coors = [4,0]
# left_up_coors = [0, 0]
# right_down_coors = [4, 4]
# middle_coors = [2, 2]
# extra_lu = [image_size/ patch_number, image_size/ patch_number]
# extra_rd = [image_size - image_size/ patch_number, image_size - image_size/ patch_number]
images, labels = generator(train_examples, image_size, rectangle_size, left_up_coors, right_down_coors, middle_coors, fix_coors, extra_lu, extra_rd)
# save_h5('dataset/full_train_random_dataset_%dx%d.h5' % (patch_number, patch_number), images = images, labels = labels)
images, labels = generator(test_examples, image_size, rectangle_size, left_up_coors, right_down_coors, middle_coors, fix_coors, extra_lu, extra_rd)
# save_h5('dataset/full_test_random_dataset_%dx%d.h5' % (patch_number, patch_number), images = images, labels = labels)
np.savez('dataset/dataset_%dx%d.npz' % (patch_number, patch_number), images = images, labels = labels)