def loadData(self):
Palette = ColorBlind_10
NMFiles = self.getFileNames(self.Args.nocs_maps)
ColorFiles = [None] * len(NMFiles)
PoseFiles = [None] * len(NMFiles)
if self.Args.colors is not None:
ColorFiles = self.getFileNames(self.Args.colors)
if self.Args.poses is not None:
PoseFiles = self.getFileNames(self.Args.poses)
for (NMF, Color, CF, PF) in zip(NMFiles, Palette.colors, ColorFiles, PoseFiles):
NOCSMap = cv2.imread(NMF, -1)
NOCSMap = NOCSMap[:, :, :3] # Ignore alpha if present
NOCSMap = cv2.cvtColor(NOCSMap, cv2.COLOR_BGR2RGB) # IMPORTANT: OpenCV loads as BGR, so convert to RGB
if self.Intrinsics is None:
self.ImageSize = (NOCSMap.shape[1], NOCSMap.shape[0])
else:
NOCSMap = self.resizeAndPad(NOCSMap)
CFIm = None
if CF is not None:
CFIm = cv2.imread(CF)
CFIm = cv2.cvtColor(CFIm, cv2.COLOR_BGR2RGB) # IMPORTANT: OpenCV loads as BGR, so convert to RGB
if CFIm.shape != NOCSMap.shape: # Re-size only if not the same size as NOCSMap
CFIm = cv2.resize(CFIm, (NOCSMap.shape[1], NOCSMap.shape[0]), interpolation=cv2.INTER_CUBIC) # Ok to use cubic interpolation for RGB
NOCS = ds.NOCSMap(NOCSMap, RGB=CFIm)
self.NOCSMaps.append(NOCSMap)
self.NOCS.append(NOCS)
if self.Args.no_pose == False:
_, K, R, C, Flip = self.estimateCameraPoseFromNM(NOCSMap, NOCS, N=self.Args.num_points, Intrinsics=self.Intrinsics) # The rotation and translation are about the NOCS origin
self.CamIntrinsics.append(K)
self.CamRots.append(R)
self.CamPos.append(C)
self.CamFlip.append(Flip)
self.Cameras.append(ds.Camera(ds.CameraExtrinsics(self.CamRots[-1], self.CamPos[-1]), ds.CameraIntrinsics(self.CamIntrinsics[-1])))
if PF is not None:
with open(PF) as f:
data = json.load(f)
# Loading convention: Flip sign of x postiion, flip signs of quaternion z, w
P = np.array([data['position']['x'], data['position']['y'], data['position']['z']])
Quat = np.array([data['rotation']['w'], data['rotation']['x'], data['rotation']['y'], data['rotation']['z']]) # NOTE: order is w, x, y, z
# Cajole transforms to work
P[0] *= -1
P += 0.5
Quat = np.array([Quat[0], Quat[1], -Quat[2], -Quat[3]])
self.PosesPos.append(P)
R = quaternions.quat2mat(Quat).T
self.PosesRots.append(R)
else:
self.PosesPos.append(None)
self.PosesRots.append(None)
self.nNM = len(NMFiles)
self.activeNMIdx = self.nNM # len(NMFiles) will show all
def NOCS_to_scaled_pc(path):
if os.path.exists(path + '/nocs_pc_scaled.npy'):
return
nocs = cv2.imread(path + '/nocs.png')
nocs = cv2.cvtColor(nocs, cv2.COLOR_BGR2RGB)
nocs = ds.NOCSMap(nocs)
Points = nocs.Points
total_size = (Points.max(axis=0) - Points.min(axis=0)).max()
centers = (Points.max(axis=0) - Points.min(axis=0)) / 2
Points -= centers
Points /= total_size
np.save(path + '/nocs_pc_scaled.npy', Points)
def generateDiffMap(self):
if self.Args.error_viz != -1:
self.isVizError = True
self.ErrorReferenceNM = self.Args.error_viz
if self.ErrorReferenceNM >= len(self.NOCSMaps):
print('[ WARN ]: Resetting reference NOCS map to 0 for error computation.')
self.ErrorReferenceNM = 0 # Reset to first
if self.isVizError == True:
for i in range(0, len(self.NOCSMaps)):
DM = self.NOCSMaps[i].astype(np.float)-self.NOCSMaps[self.ErrorReferenceNM].astype(np.float) # Convert to float
Norm = np.linalg.norm(DM, axis=2)
Frac = 10
NormFact = (441.6729 / Frac) / 255 # Maximum possible error in NOCS is 441.6729 == sqrt(3 * 255^2). Let's take a fraction of that
Norm = Norm / (NormFact)
Norm = Norm.astype(np.uint8)
NormCol = cv2.applyColorMap(Norm, cv2.COLORMAP_JET)
#cv2.imwrite('norm_{}.png'.format(str(i).zfill(3)), NormCol)
self.NOCS[i] = ds.NOCSMap(self.NOCSMaps[i], RGB=cv2.cvtColor(NormCol, cv2.COLOR_BGR2RGB))# IMPORTANT: OpenCV loads as BGR, so convert to RGB
def log_batch(self, batch):
# get data
if not self.NAME in batch['parser'].keys():
return
keys_list = batch['parser'][self.NAME]
if len(keys_list) == 0:
return
data = batch['data']
phase = batch['phase']
current_batch = batch['batch-id']
current_epoch = batch['epoch-id']
meta_info = batch['meta-info'] # {k:[info0,info1,info2,....]}
# check dir
os.makedirs(os.path.join(self.log_path, 'epoch_%d' % current_epoch),
exist_ok=True)
# check whether need log
if phase.lower() == 'train':
if current_batch % self.visual_train_interval_batch != 0 or current_batch % self.visual_interval_batch != 0:
return
else:
if current_epoch % self.visual_interval_epoch != 0 or current_batch % self.visual_interval_batch != 0:
return
# for each key
for obj_key in keys_list: # for each key
nocs = data[obj_key[0]]
color = data[obj_key[1]]
if isinstance(nocs, list):
assert isinstance(color, list)
assert len(nocs[0].shape) == 4
assert len(color[0].shape) == 4
assert nocs[0].shape[0] == color[0].shape[0]
nbatch = nocs[0].shape[0]
else:
assert len(nocs.shape) == 4
assert len(color.shape) == 4
assert nocs.shape[0] == color.shape[0]
nbatch = nocs.shape[0]
nocs = [nocs]
color = [color]
# convert to ndarray
if isinstance(nocs[0], torch.Tensor):
for i, tensor in enumerate(nocs):
nocs[i] = tensor.detach().cpu().numpy()
if isinstance(color[0], torch.Tensor):
for i, tensor in enumerate(color):
color[i] = tensor.detach().cpu().numpy()
nocs = deepcopy(nocs)
color = deepcopy(color)
# for each sample in batch
for batch_id in range(nbatch):
# get meta postfix
prefix = ""
for k, v in meta_info.items():
prefix += k + "_" + str(v[batch_id]) + "_"
# now all case convert to list of image
nview = len(nocs)
for view_id in range(nview):
nocs_map = nocs[view_id][batch_id] # 3*H*W [0,1]
color_map = color[view_id][batch_id] # 3*H*W [0,1]
assert nocs_map.ndim == 3
assert nocs_map.shape[0] == 3
assert color_map.ndim == 3
assert color_map.shape[0] == 3
nocs_map = nocs_map.transpose(1, 2, 0) # H*W*3
color_map = color_map.transpose(1, 2, 0) # H*W*3
# smaller meshes
nocs_map = cv.resize(nocs_map,
dsize=(int(nocs_map.shape[1] / 4),
int(nocs_map.shape[0] / 4)),
interpolation=cv.INTER_NEAREST)
color_map = cv.resize(color_map,
dsize=(int(color_map.shape[1] / 4),
int(color_map.shape[0] / 4)),
interpolation=cv.INTER_NEAREST)
# save to file
filename = os.path.join(
self.log_path, 'epoch_%d' % current_epoch,
prefix + 'b%d_v%d_' % (batch_id, view_id) +
obj_key[0] + '-' + obj_key[1] + '.obj')
# save here
# need to convert color, need to *255
if nocs_map.max() < 1 + 1e-4:
nocs_map = nocs_map[:, :, [2, 1, 0]] * 255.0
if color_map.max() < 1 + 1e-4:
color_map = color_map[:, :, [2, 1, 0]] * 255.0
tk3dv_nocs_mp = ds.NOCSMap(nocs_map, RGB=color_map)
tk3dv_nocs_mp.serialize(filename)
if self.visual_one:
break