Skip to content

Commit a1b92c7

Browse files
author
csguestp
committed
first commit
0 parents  commit a1b92c7

16 files changed

+5317
-0
lines changed

‎datasets/scannet_dataset.py

+246
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,246 @@
1+
from torch.utils.data import Dataset
2+
import torch
3+
import numpy as np
4+
import time
5+
import os
6+
import scipy.ndimage
7+
from utils import remapper
8+
9+
blur0=np.ones((3,1,1)).astype('float32')/3
10+
blur1=np.ones((1,3,1)).astype('float32')/3
11+
blur2=np.ones((1,1,3)).astype('float32')/3
12+
13+
def elastic(x,gran,mag):
14+
bb=np.abs(x).max(0).astype(np.int32)//gran+3
15+
noise=[np.random.randn(bb[0],bb[1],bb[2]).astype('float32') for _ in range(3)]
16+
noise=[scipy.ndimage.filters.convolve(n,blur0,mode='constant',cval=0) for n in noise]
17+
noise=[scipy.ndimage.filters.convolve(n,blur1,mode='constant',cval=0) for n in noise]
18+
noise=[scipy.ndimage.filters.convolve(n,blur2,mode='constant',cval=0) for n in noise]
19+
noise=[scipy.ndimage.filters.convolve(n,blur0,mode='constant',cval=0) for n in noise]
20+
noise=[scipy.ndimage.filters.convolve(n,blur1,mode='constant',cval=0) for n in noise]
21+
noise=[scipy.ndimage.filters.convolve(n,blur2,mode='constant',cval=0) for n in noise]
22+
ax=[np.linspace(-(b-1)*gran,(b-1)*gran,b) for b in bb]
23+
24+
interp=[scipy.interpolate.RegularGridInterpolator(ax,n,bounds_error=0,fill_value=0) for n in noise]
25+
def g(x_):
26+
return np.hstack([i(x_)[:,None] for i in interp])
27+
noise = g(x)
28+
return x+g(x)*mag
29+
30+
## ScanNet dataset class
31+
class ScanNetDataset(Dataset):
32+
def __init__(self, options, split, random=True):
33+
self.options = options
34+
self.split = split
35+
self.random = random
36+
self.imagePaths = []
37+
self.dataFolder = '/gruvi/Data/chenliu/ScanNet/scans/'
38+
39+
with open('split_' + split + '.txt', 'r') as f:
40+
for line in f:
41+
scene_id = line.strip()
42+
if len(scene_id) < 5 or scene_id[:5] != 'scene':
43+
continue
44+
if options.scene_id != '' and options.scene_id not in scene_id:
45+
continue
46+
filename = self.dataFolder + '/' + scene_id + '/' + scene_id + '_vh_clean_2.pth'
47+
if os.path.exists(filename):
48+
info = torch.load(filename)
49+
if len(info) == 5:
50+
self.imagePaths.append(filename)
51+
52+
#np.savetxt('semantic_val/' + scene_id + '.txt', info[2], fmt='%d')
53+
pass
54+
pass
55+
if split != 'train' and len(self.imagePaths) >= options.numTestingImages:
56+
break
57+
continue
58+
pass
59+
60+
#self.imagePaths = [filename for filename in self.imagePaths if 'scene0217_00' in filename]
61+
62+
if options.numTrainingImages > 0 and split == 'train':
63+
self.numImages = options.numTrainingImages
64+
else:
65+
self.numImages = len(self.imagePaths)
66+
pass
67+
return
68+
69+
def __len__(self):
70+
return self.numImages
71+
72+
def __getitem__(self, index):
73+
if self.random:
74+
t = int(time.time() * 1000000)
75+
np.random.seed(((t & 0xff000000) >> 24) +
76+
((t & 0x00ff0000) >> 8) +
77+
((t & 0x0000ff00) << 8) +
78+
((t & 0x000000ff) << 24))
79+
index = np.random.randint(len(self.imagePaths))
80+
else:
81+
index = index % len(self.imagePaths)
82+
pass
83+
84+
debug = -1
85+
if debug >= 0:
86+
index = debug
87+
print(index, self.imagePaths[index])
88+
pass
89+
90+
coords, colors, labels, instances, faces = torch.load(self.imagePaths[index])
91+
invalid_instances, = torch.load(self.imagePaths[index].replace('.pth', '_invalid.pth'))
92+
93+
labels = remapper[labels]
94+
95+
#neighbor_gt = torch.load(self.imagePaths[index].replace('.pth', '_neighbor.pth'))
96+
#print(neighbor_gt[0])
97+
#exit(1)
98+
#neighbor_gt = 1
99+
#print(coords.min(0), coords.max(0))
100+
if self.split == 'train':
101+
m = np.eye(3) + np.random.randn(3,3) * 0.1
102+
m[0][0] *= np.random.randint(2) * 2 - 1
103+
theta = np.random.rand() * 2 * np.pi
104+
else:
105+
m = np.eye(3)
106+
theta = 0
107+
pass
108+
109+
scale = self.options.scanScale
110+
full_scale = self.options.inputScale
111+
m *= scale
112+
m = np.matmul(m, [[np.cos(theta), np.sin(theta),0], [-np.sin(theta), np.cos(theta),0], [0,0,1]])
113+
coords = np.matmul(coords, m)
114+
if self.split == 'train':
115+
coords = elastic(coords, 6 * scale // 50,40 * scale / 50)
116+
#coords = elastic(coords, 20 * scale // 50, 160 * scale / 50)
117+
pass
118+
119+
if 'normal' in self.options.suffix:
120+
points_1 = coords[faces[:, 0]]
121+
points_2 = coords[faces[:, 1]]
122+
points_3 = coords[faces[:, 2]]
123+
face_normals = np.cross(points_2 - points_1, points_3 - points_1)
124+
face_normals /= np.maximum(np.linalg.norm(face_normals, axis=-1, keepdims=True), 1e-4)
125+
normals = np.zeros((len(coords), 3))
126+
for c in range(3):
127+
np.add.at(normals, faces[:, c], face_normals)
128+
continue
129+
normals /= np.maximum(np.linalg.norm(normals, axis=-1, keepdims=True), 1e-4)
130+
colors = np.concatenate([colors, normals], axis=-1)
131+
pass
132+
133+
if self.split == 'train':
134+
colors[:, :3] = colors[:, :3] + np.random.randn(3) * 0.1
135+
pass
136+
137+
if self.options.trainingMode == 'semantic':
138+
unique_instances, indices, instances = np.unique(instances, return_index=True, return_inverse=True)
139+
labels = labels[indices]
140+
labels[labels == -100] = 20
141+
new_coords = np.zeros(coords.shape, dtype=coords.dtype)
142+
for instance in range(len(unique_instances)):
143+
instance_mask = instances == instance
144+
instance_coords = coords[instance_mask]
145+
mins = instance_coords.min(0)
146+
maxs = instance_coords.max(0)
147+
max_range = (maxs - mins).max()
148+
if self.split == 'train':
149+
padding = (maxs - mins) * np.random.random(3) * 0.1
150+
else:
151+
padding = max_range * 0.05
152+
pass
153+
max_range += padding * 2
154+
mins = (mins + maxs) / 2 - max_range / 2
155+
instance_coords = np.clip(np.round((instance_coords - mins) / max_range * full_scale), 0, full_scale - 1)
156+
new_coords[instance_mask] = instance_coords
157+
continue
158+
coords = np.concatenate([new_coords, np.expand_dims(instances, -1)], axis=-1)
159+
sample = [coords.astype(np.int64), colors.astype(np.float32), faces.astype(np.int64), labels.astype(np.int64), instances.astype(np.int64), self.imagePaths[index]]
160+
return sample
161+
162+
if self.options.trainingMode == 'confidence':
163+
scene_id = self.imagePaths[index].split('/')[-1].split('_vh_clean_2')[0]
164+
info = torch.load('test/output_normal_augment_2_' + self.split + '/cache/' + scene_id + '.pth')
165+
if len(info) == 2:
166+
semantic_pred, instance_pred = info
167+
else:
168+
semantic_pred, instance_pred = info[3], info[6]
169+
semantic_pred = semantic_pred[:len(coords)]
170+
instance_pred = instance_pred[:len(coords)]
171+
pass
172+
instance_pred += 1
173+
unique_instances, indices, counts = np.unique(instances, return_index=True, return_counts=True)
174+
175+
instance_counts = np.zeros(unique_instances.max() + 1)
176+
instance_counts[unique_instances] = counts
177+
instance_semantics = np.zeros(unique_instances.max() + 1)
178+
instance_semantics[unique_instances] = labels[indices]
179+
confidence_gt = []
180+
semantic_gt = []
181+
instance_masks = []
182+
new_coords = np.zeros(coords.shape, dtype=coords.dtype)
183+
for instance in range(instance_pred.max() + 1):
184+
instance_mask = instance_pred == instance
185+
if instance_mask.sum() == 0:
186+
print('sum = 0', instance, instance_pred.max() + 1, instance_mask.sum())
187+
exit(1)
188+
info = np.unique(semantic_pred[instance_mask > 0.5], return_counts=True)
189+
label_pred = info[0][info[1].argmax()]
190+
info = np.unique(instances[instance_mask > 0.5], return_counts=True)
191+
instance_gt = info[0][info[1].argmax()]
192+
193+
instance_coords = coords[instance_mask]
194+
mins = instance_coords.min(0)
195+
maxs = instance_coords.max(0)
196+
max_range = (maxs - mins).max()
197+
if self.split == 'train':
198+
padding = (maxs - mins) * np.random.random(3) * 0.1
199+
else:
200+
padding = max_range * 0.05
201+
pass
202+
max_range += padding * 2
203+
mins = (mins + maxs) / 2 - max_range / 2
204+
instance_coords = np.clip(np.round((instance_coords - mins) / max_range * full_scale), 0, full_scale - 1)
205+
new_coords[instance_mask] = instance_coords
206+
207+
if instance > 0:
208+
confidence_gt.append(int(label_pred == instance_semantics[instance_gt] and info[1].max() > 0.5 * instance_counts[instance_gt]))
209+
semantic_gt.append(label_pred)
210+
instance_masks.append(instance_mask)
211+
pass
212+
continue
213+
coords = np.concatenate([new_coords, np.expand_dims(instance_pred, -1)], axis=-1)
214+
sample = [coords.astype(np.int64), colors.astype(np.float32), faces.astype(np.int64), np.stack(semantic_gt).astype(np.int64), np.stack(confidence_gt).astype(np.int64), self.imagePaths[index], np.stack(instance_masks).astype(np.int32)]
215+
return sample
216+
217+
218+
mins = coords.min(0)
219+
maxs = coords.max(0)
220+
#ranges = maxs - mins
221+
if self.split == 'train':
222+
offset = -mins + np.clip(full_scale - maxs + mins - 0.001, 0, None) * np.random.rand(3) + np.clip(full_scale - maxs + mins + 0.001, None, 0) * np.random.rand(3)
223+
coords += offset
224+
else:
225+
coords -= (mins + maxs) // 2 - full_scale // 2
226+
#coords -= mins
227+
pass
228+
229+
coords = np.round(coords)
230+
if False:
231+
idxs = (coords.min(1) >= 0) * (coords.max(1) < full_scale)
232+
coords = coords[idxs]
233+
colors = colors[idxs]
234+
labels = labels[idxs]
235+
instances = instances[idxs]
236+
invalid_instances = invalid_instances[idxs]
237+
else:
238+
#print(coords.min(0), coords.max(0))
239+
#exit(1)
240+
coords = np.clip(coords, 0, full_scale - 1)
241+
pass
242+
243+
coords = np.concatenate([coords, np.full((coords.shape[0], 1), fill_value=index)], axis=-1)
244+
#coords = np.concatenate([coords, np.expand_dims(instances, -1)], axis=-1)
245+
sample = [coords.astype(np.int64), colors.astype(np.float32), faces.astype(np.int64), labels.astype(np.int64), instances.astype(np.int64), invalid_instances.astype(np.int64), self.imagePaths[index]]
246+
return sample

0 commit comments

Comments
 (0)