-
Notifications
You must be signed in to change notification settings - Fork 1
/
dataset.py
60 lines (44 loc) · 1.58 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
'''Reads and parses training data and labels'''
import os
import os.path
import binvox_rw
import numpy as np
from PIL import Image
import tensorflow as tf
TEST_DATA = "03001627" #03001627
MAIN_DATA = "03211117"
BATCH_SIZE = 2
label_dir = os.listdir('./'+TEST_DATA+'_labels')
data = os.listdir('./'+TEST_DATA+'')
TOTAL_SIZE = len(data)
def train_labels():
if(len(label_dir)<BATCH_SIZE):
return []
y_train = {}
#d = {} # Keys = IDs for items
# Values = binvox
for _ in range(BATCH_SIZE):
label = label_dir.pop()
if(label.startswith('.')):
continue
binv = open('./'+TEST_DATA+'_labels/'+label+'/model.binvox','rb')
binvox_data = binvox_rw.read_as_3d_array(binv).data # binvox_data is 32x32x32
y_train[label] = np.asarray(binvox_data)
return y_train
def train_data():
if(len(data)<BATCH_SIZE):
return []
x_train = {} #Keys = IDs for items, Values = 24 pictures
for _ in range(BATCH_SIZE):
item = data.pop()
tmp_im_array = []
if(item.startswith('.')):
continue
for pic in os.listdir('./'+TEST_DATA+'/'+item+'/rendering'):
if('.png' in pic): #If current item is a picture, store it
im = np.array(Image.open('./'+TEST_DATA+'/'+item+'/rendering/'+pic))
im = tf.random_crop(im,[127,127,3]) # Input images are [137,137,3]. Remove random 10x10 pixel area.
tmp_im_array.append(im)
x_train[item] = tmp_im_array
tmp_im_array = []
return x_train