-
Notifications
You must be signed in to change notification settings - Fork 1
/
data.py
88 lines (85 loc) · 3.39 KB
/
data.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
import numpy as np
import scipy.linalg
import scipy.io
import os,time
import threading
import tensorflow as tf
import util
class Loader():
def __init__(self,opt,loadNovel=True,loadFixedOut=False,loadTest=False):
self.loadNovel = loadNovel
self.loadFixedOut = loadFixedOut
self.loadTest = loadTest
listFile = "data/{0}_{1}.list".format(opt.category,"test" if loadTest else "train")
self.CADs = []
with open(listFile) as file:
for line in file:
id = line.strip().split("/")[1]
self.CADs.append(id)
self.CADs.sort()
def loadChunk(self,opt,loadRange=None):
data = {}
if loadRange is not None: idx = np.arange(loadRange[0],loadRange[1])
else: idx = np.random.permutation(len(self.CADs))[:opt.chunkSize]
chunkSize = len(idx)
# preallocate memory
data["image_in"] = np.ones([chunkSize,24,opt.inH,opt.inW,3],dtype=np.float32)
if self.loadNovel:
data["depth"] = np.ones([chunkSize,opt.sampleN,opt.H,opt.W],dtype=np.float32)
data["mask"] = np.ones([chunkSize,opt.sampleN,opt.H,opt.W],dtype=np.bool)
data["trans"] = np.ones([chunkSize,opt.sampleN,4],dtype=np.float32)
if self.loadFixedOut:
data["depth_fixedOut"] = np.ones([chunkSize,opt.outViewN,opt.outH,opt.outW],dtype=np.float32)
data["mask_fixedOut"] = np.ones([chunkSize,opt.outViewN,opt.outH,opt.outW],dtype=np.bool)
# load data
for c in range(chunkSize):
CAD = self.CADs[idx[c]]
data["image_in"][c] = np.load("data/{0}_inputRGB/{1}.npy".format(opt.category,CAD))/255.0
if self.loadNovel:
rawData = scipy.io.loadmat("data/{0}_depth/{1}.mat".format(opt.category,CAD))
depth = rawData["Z"]
trans = rawData["trans"]
mask = depth!=0
depth[~mask] = opt.renderDepth
# store data
data["depth"][c] = depth
data["mask"][c] = mask
data["trans"][c] = trans
if self.loadFixedOut:
rawData_fixed = scipy.io.loadmat("data/{0}_depth_fixed{1}/{2}.mat".format(opt.category,opt.outViewN,CAD))
depth_fixed = rawData_fixed["Z"]
mask_fixed = depth_fixed!=0
depth_fixed[~mask_fixed] = opt.renderDepth
# store data
data["depth_fixedOut"][c] = depth_fixed
data["mask_fixedOut"][c] = mask_fixed
self.pendingChunk = data
def shipChunk(self):
self.readyChunk,self.pendingChunk = self.pendingChunk,None
# make training batch
def makeBatch(opt,dataloader,PH):
data = dataloader.readyChunk
inputImage,targetTrans,depthGT,maskGT = PH
modelIdx = np.random.permutation(opt.chunkSize)[:opt.batchSize]
modelIdxTile = np.tile(modelIdx,[opt.novelN,1]).T
angleIdx = np.random.randint(24,size=[opt.batchSize])
sampleIdx = np.random.randint(opt.sampleN,size=[opt.batchSize,opt.novelN])
batch = {
inputImage: data["image_in"][modelIdx,angleIdx],
targetTrans: data["trans"][modelIdxTile,sampleIdx],
depthGT: np.expand_dims(data["depth"][modelIdxTile,sampleIdx],axis=-1),
maskGT: np.expand_dims(data["mask"][modelIdxTile,sampleIdx],axis=-1)
}
return batch
# make training batch
def makeBatchFixed(opt,dataloader,PH):
data = dataloader.readyChunk
inputImage,depthGT,maskGT = PH
modelIdx = np.random.permutation(opt.chunkSize)[:opt.batchSize]
angleIdx = np.random.randint(24,size=[opt.batchSize])
batch = {
inputImage: data["image_in"][modelIdx,angleIdx],
depthGT: np.transpose(data["depth_fixedOut"][modelIdx],axes=[0,2,3,1]),
maskGT: np.transpose(data["mask_fixedOut"][modelIdx],axes=[0,2,3,1])
}
return batch