def dump_results(args, scanrefer, data, config):
dump_dir = os.path.join(CONF.PATH.BASE, args.folder, "vis")
os.makedirs(dump_dir, exist_ok=True)
# from inputs
ids = data['scan_idx'].detach().cpu().numpy()
point_clouds = data['point_clouds'].cpu().numpy()
test_point_clouds = data["test_point_clouds"].cpu().numpy() # one object
batch_size = point_clouds.shape[0]
pcl_color = data["pcl_color"].detach().cpu().numpy()
test_pcl_color = data["test_pcl_color"].detach().cpu().numpy()
if args.use_color:
pcl_color = (pcl_color * 256 + MEAN_COLOR_RGB).astype(np.int64)
test_pcl_color = (test_pcl_color * 256 + MEAN_COLOR_RGB).astype(
np.int64)
for i in range(batch_size):
# basic info
idx = ids[i]
scene_id = scanrefer[idx]["scene_id"]
object_id = scanrefer[idx]["object_id"]
object_name = scanrefer[idx]["object_name"]
ann_id = scanrefer[idx]["ann_id"]
# scene_output
scene_dump_dir = os.path.join(dump_dir, scene_id)
# if not os.path.exists(scene_dump_dir):
os.makedirs(scene_dump_dir, exist_ok=True)
# # Dump the original scene point clouds
# mesh = align_mesh(scene_id)
# mesh.write(os.path.join(scene_dump_dir, 'mesh.ply'))
print("pcl_color[i] in pc.ply", pcl_color[i])
point_clouds = point_clouds[i] #[:200]
pcl_color = pcl_color[i] #[:200]
test_point_clouds = test_point_clouds[i]
test_pcl_color = test_pcl_color[i]
jitter_idx = random.random() * 0.45 + 0.8
write_ply_rgb(point_clouds, pcl_color,
os.path.join(scene_dump_dir, 'pc_full.ply'))
write_ply_rgb(test_point_clouds, test_pcl_color,
os.path.join(scene_dump_dir, 'pc_full_obj.ply'))
write_ply_rgb(test_point_clouds * jitter_idx, test_pcl_color,
os.path.join(scene_dump_dir, 'pc_full_obj2.ply'))
# write_ply_rgb(point_clouds[i], pcl_color[i], os.path.join(scene_dump_dir, 'pc.ply'))
break
def dump_results(args, scanrefer, data, config):
dump_dir = os.path.join(CONF.PATH.OUTPUT, args.folder, "vis")
os.makedirs(dump_dir, exist_ok=True)
# from inputs
ids = data['scan_idx'].detach().cpu().numpy()
point_clouds = data['point_clouds'].cpu().numpy()
batch_size = point_clouds.shape[0]
pcl_color = data["pcl_color"].detach().cpu().numpy()
if args.use_color:
pcl_color = (pcl_color * 256 + MEAN_COLOR_RGB).astype(np.int64)
# from network outputs
# detection
pred_objectness = torch.argmax(data['objectness_scores'],
2).float().detach().cpu().numpy()
pred_center = data['center'].detach().cpu().numpy() # (B,K,3)
pred_heading_class = torch.argmax(data['heading_scores'],
-1) # B,num_proposal
pred_heading_residual = torch.gather(
data['heading_residuals'], 2,
pred_heading_class.unsqueeze(-1)) # B,num_proposal,1
pred_heading_class = pred_heading_class.detach().cpu().numpy(
) # B,num_proposal
pred_heading_residual = pred_heading_residual.squeeze(
2).detach().cpu().numpy() # B,num_proposal
pred_size_class = torch.argmax(data['size_scores'], -1) # B,num_proposal
pred_size_residual = torch.gather(
data['size_residuals'], 2,
pred_size_class.unsqueeze(-1).unsqueeze(-1).repeat(
1, 1, 1, 3)) # B,num_proposal,1,3
pred_size_residual = pred_size_residual.squeeze(
2).detach().cpu().numpy() # B,num_proposal,3
# reference
pred_ref_scores = data["cluster_ref"].detach().cpu().numpy()
pred_ref_scores_softmax = F.softmax(
data["cluster_ref"] *
torch.argmax(data['objectness_scores'], 2).float() * data['pred_mask'],
dim=1).detach().cpu().numpy()
# post-processing
nms_masks = data['pred_mask'].detach().cpu().numpy() # B,num_proposal
# ground truth
gt_center = data['center_label'].cpu().numpy() # (B,MAX_NUM_OBJ,3)
gt_heading_class = data['heading_class_label'].cpu().numpy() # B,K2
gt_heading_residual = data['heading_residual_label'].cpu().numpy() # B,K2
gt_size_class = data['size_class_label'].cpu().numpy() # B,K2
gt_size_residual = data['size_residual_label'].cpu().numpy() # B,K2,3
# reference
gt_ref_labels = data["ref_box_label"].detach().cpu().numpy()
for i in range(batch_size):
# basic info
idx = ids[i]
scene_id = scanrefer[idx]["scene_id"]
object_id = scanrefer[idx]["object_id"]
object_name = scanrefer[idx]["object_name"]
ann_id = scanrefer[idx]["ann_id"]
# scene_output
scene_dump_dir = os.path.join(dump_dir, scene_id)
if not os.path.exists(scene_dump_dir):
os.mkdir(scene_dump_dir)
# # Dump the original scene point clouds
mesh = align_mesh(scene_id)
mesh.write(os.path.join(scene_dump_dir, 'mesh.ply'))
write_ply_rgb(point_clouds[i], pcl_color[i],
os.path.join(scene_dump_dir, 'pc.ply'))
# filter out the valid ground truth reference box
assert gt_ref_labels[i].shape[0] == gt_center[i].shape[0]
gt_ref_idx = np.argmax(gt_ref_labels[i], 0)
# visualize the gt reference box
# NOTE: for each object there should be only one gt reference box
object_dump_dir = os.path.join(
dump_dir, scene_id, "gt_{}_{}.ply".format(object_id, object_name))
gt_obb = config.param2obb(gt_center[i, gt_ref_idx,
0:3], gt_heading_class[i,
gt_ref_idx],
gt_heading_residual[i, gt_ref_idx],
gt_size_class[i, gt_ref_idx],
gt_size_residual[i, gt_ref_idx])
gt_bbox = get_3d_box(gt_obb[3:6], gt_obb[6], gt_obb[0:3])
if not os.path.exists(object_dump_dir):
write_bbox(
gt_obb, 0,
os.path.join(scene_dump_dir,
'gt_{}_{}.ply'.format(object_id, object_name)))
# find the valid reference prediction
pred_masks = nms_masks[i] * pred_objectness[i] == 1
assert pred_ref_scores[i].shape[0] == pred_center[i].shape[0]
pred_ref_idx = np.argmax(pred_ref_scores[i] * pred_masks, 0)
assigned_gt = torch.gather(
data["ref_box_label"], 1,
data["object_assignment"]).detach().cpu().numpy()
# visualize the predicted reference box
pred_obb = config.param2obb(pred_center[i, pred_ref_idx, 0:3],
pred_heading_class[i, pred_ref_idx],
pred_heading_residual[i, pred_ref_idx],
pred_size_class[i, pred_ref_idx],
pred_size_residual[i, pred_ref_idx])
pred_bbox = get_3d_box(pred_obb[3:6], pred_obb[6], pred_obb[0:3])
iou = box3d_iou(gt_bbox, pred_bbox)
write_bbox(
pred_obb, 1,
os.path.join(
scene_dump_dir, 'pred_{}_{}_{}_{:.5f}_{:.5f}.ply'.format(
object_id, object_name, ann_id,
pred_ref_scores_softmax[i, pred_ref_idx], iou)))
def write_bbox(bbox, mode, output_file):
"""
bbox: (cx, cy, cz, lx, ly, lz, r), center and length in three axis, the last is the rotation
output_file: string
"""
def create_cylinder_mesh(radius, p0, p1, stacks=10, slices=10):
import math
def compute_length_vec3(vec3):
return math.sqrt(vec3[0]*vec3[0] + vec3[1]*vec3[1] + vec3[2]*vec3[2])
def rotation(axis, angle):
rot = np.eye(4)
c = np.cos(-angle)
s = np.sin(-angle)
t = 1.0 - c
axis /= compute_length_vec3(axis)
x = axis[0]
y = axis[1]
z = axis[2]
rot[0,0] = 1 + t*(x*x-1)
rot[0,1] = z*s+t*x*y
rot[0,2] = -y*s+t*x*z
rot[1,0] = -z*s+t*x*y
rot[1,1] = 1+t*(y*y-1)
rot[1,2] = x*s+t*y*z
rot[2,0] = y*s+t*x*z
rot[2,1] = -x*s+t*y*z
rot[2,2] = 1+t*(z*z-1)
return rot
verts = []
indices = []
diff = (p1 - p0).astype(np.float32)
height = compute_length_vec3(diff)
for i in range(stacks+1):
for i2 in range(slices):
theta = i2 * 2.0 * math.pi / slices
pos = np.array([radius*math.cos(theta), radius*math.sin(theta), height*i/stacks])
verts.append(pos)
for i in range(stacks):
for i2 in range(slices):
i2p1 = math.fmod(i2 + 1, slices)
indices.append( np.array([(i + 1)*slices + i2, i*slices + i2, i*slices + i2p1], dtype=np.uint32) )
indices.append( np.array([(i + 1)*slices + i2, i*slices + i2p1, (i + 1)*slices + i2p1], dtype=np.uint32) )
transform = np.eye(4)
va = np.array([0, 0, 1], dtype=np.float32)
vb = diff
vb /= compute_length_vec3(vb)
axis = np.cross(vb, va)
angle = np.arccos(np.clip(np.dot(va, vb), -1, 1))
if angle != 0:
if compute_length_vec3(axis) == 0:
dotx = va[0]
if (math.fabs(dotx) != 1.0):
axis = np.array([1,0,0]) - dotx * va
else:
axis = np.array([0,1,0]) - va[1] * va
axis /= compute_length_vec3(axis)
transform = rotation(axis, -angle)
transform[:3,3] += p0
verts = [np.dot(transform, np.array([v[0], v[1], v[2], 1.0])) for v in verts]
verts = [np.array([v[0], v[1], v[2]]) / v[3] for v in verts]
return verts, indices
def get_bbox_edges(bbox_min, bbox_max):
def get_bbox_verts(bbox_min, bbox_max):
verts = [
np.array([bbox_min[0], bbox_min[1], bbox_min[2]]),
np.array([bbox_max[0], bbox_min[1], bbox_min[2]]),
np.array([bbox_max[0], bbox_max[1], bbox_min[2]]),
np.array([bbox_min[0], bbox_max[1], bbox_min[2]]),
np.array([bbox_min[0], bbox_min[1], bbox_max[2]]),
np.array([bbox_max[0], bbox_min[1], bbox_max[2]]),
np.array([bbox_max[0], bbox_max[1], bbox_max[2]]),
np.array([bbox_min[0], bbox_max[1], bbox_max[2]])
]
return verts
box_verts = get_bbox_verts(bbox_min, bbox_max)
edges = [
(box_verts[0], box_verts[1]),
(box_verts[1], box_verts[2]),
(box_verts[2], box_verts[3]),
(box_verts[3], box_verts[0]),
(box_verts[4], box_verts[5]),
(box_verts[5], box_verts[6]),
(box_verts[6], box_verts[7]),
(box_verts[7], box_verts[4]),
(box_verts[0], box_verts[4]),
(box_verts[1], box_verts[5]),
(box_verts[2], box_verts[6]),
(box_verts[3], box_verts[7])
]
return edges
def get_bbox_corners(bbox):
centers, lengths = bbox[:3], bbox[3:6]
xmin, xmax = centers[0] - lengths[0] / 2, centers[0] + lengths[0] / 2
ymin, ymax = centers[1] - lengths[1] / 2, centers[1] + lengths[1] / 2
zmin, zmax = centers[2] - lengths[2] / 2, centers[2] + lengths[2] / 2
corners = []
corners.append(np.array([xmax, ymax, zmax]).reshape(1, 3))
corners.append(np.array([xmax, ymax, zmin]).reshape(1, 3))
corners.append(np.array([xmin, ymax, zmin]).reshape(1, 3))
corners.append(np.array([xmin, ymax, zmax]).reshape(1, 3))
corners.append(np.array([xmax, ymin, zmax]).reshape(1, 3))
corners.append(np.array([xmax, ymin, zmin]).reshape(1, 3))
corners.append(np.array([xmin, ymin, zmin]).reshape(1, 3))
corners.append(np.array([xmin, ymin, zmax]).reshape(1, 3))
corners = np.concatenate(corners, axis=0) # 8 x 3
return corners
radius = 0.03
offset = [0,0,0]
verts = []
indices = []
colors = []
corners = get_bbox_corners(bbox)
box_min = np.min(corners, axis=0)
box_max = np.max(corners, axis=0)
palette = {
0: [0, 255, 0], # gt
1: [0, 0, 255] # pred
}
chosen_color = palette[mode]
edges = get_bbox_edges(box_min, box_max)
for k in range(len(edges)):
cyl_verts, cyl_ind = create_cylinder_mesh(radius, edges[k][0], edges[k][1])
cur_num_verts = len(verts)
cyl_color = [[c / 255 for c in chosen_color] for _ in cyl_verts]
cyl_verts = [x + offset for x in cyl_verts]
cyl_ind = [x + cur_num_verts for x in cyl_ind]
verts.extend(cyl_verts)
indices.extend(cyl_ind)
colors.extend(cyl_color)
# print("in scenes color", colors)
write_ply_rgb(np.array(verts), np.array(colors), output_file)