in meshrcnn/utils/vis.py [0:0]
def visualize_minibatch(images, data, output_dir, vis_fg=False):
import matplotlib.pyplot as plt
# create vis_dir
output_dir = os.path.join(output_dir, "minibatch")
os.makedirs(output_dir, exist_ok=True)
# read image and add mean
ims = images.tensor.cpu().numpy()
num = ims.shape[0]
pixel_mean = np.expand_dims(np.expand_dims(np.array([102.9801, 115.9465, 122.7717]), 1), 2)
proposals = data["proposals"]
if vis_fg:
proposals = data["fg_proposals"]
assert len(proposals) == num
index = np.array([prop.proposal_boxes.tensor.shape[0] for prop in proposals]).cumsum()
index = np.concatenate((np.array([0]), index[:-1]))
if "target_meshes" in data:
target_sampled_verts = sample_points_from_meshes(
data["target_meshes"], num_samples=10000, return_normals=False
)
for i in range(num):
im = ims[i, :, :, :] + pixel_mean
im = im.transpose(1, 2, 0)[:, :, (2, 1, 0)]
im = im.astype(np.uint8, copy=False)
boxes = proposals[i].proposal_boxes.tensor.cpu().numpy()
gt_classes = proposals[i].gt_classes.cpu().numpy()
for j in range(boxes.shape[0]):
fig = plt.figure()
plt.subplot(2, 3, 1)
plt.imshow(im)
roi = boxes[j]
plt.gca().add_patch(
plt.Rectangle(
(roi[0], roi[1]),
roi[2] - roi[0],
roi[3] - roi[1],
fill=False,
edgecolor="r",
linewidth=3,
)
)
plt.title("Class %d" % (gt_classes[j]))
if vis_fg and "target_masks" in data:
mask = proposals[i].gt_masks[j].cpu().numpy()
plt.subplot(2, 3, 2)
plt.imshow(mask)
plt.gca().add_patch(
plt.Rectangle(
(roi[0], roi[1]),
roi[2] - roi[0],
roi[3] - roi[1],
fill=False,
edgecolor="g",
linewidth=3,
)
)
mask = data["target_masks"][index[i] + j].cpu().numpy()
plt.subplot(2, 3, 3)
plt.imshow(mask)
plt.title("Mask")
if vis_fg and "target_voxels" in data:
voxel = data["target_voxels"][index[i] + j].cpu().numpy() # (D, H, W)
voxel = voxel.transpose(1, 2, 0)
plt.subplot(2, 3, 4)
plt.imshow(np.max(voxel, 2))
plt.title("Voxel")
if vis_fg and "target_meshes" in data:
resolution = 28
verts = target_sampled_verts[int(index[i] + j)]
verts = verts.cpu().numpy()
x = (verts[:, 0] + 1) * (resolution - 1) / 2.0
y = (verts[:, 1] + 1) * (resolution - 1) / 2.0
x_keep = np.logical_and(x >= 0, x < resolution)
y_keep = np.logical_and(y >= 0, y < resolution)
keep = np.logical_and(x_keep, y_keep)
x = x[keep].astype(np.int32)
y = y[keep].astype(np.int32)
target_img = np.zeros((resolution, resolution), dtype=np.uint8)
target_img[y, x] = 255
plt.subplot(2, 3, 5)
plt.imshow(target_img)
plt.title("Mesh")
if vis_fg and "init_meshes" in data:
resolution = 28
verts, faces = data["init_meshes"].get_mesh_verts_faces(int(index[i] + j))
verts = verts.cpu().numpy()
x = (verts[:, 0] + 1) * (resolution - 1) / 2.0
y = (verts[:, 1] + 1) * (resolution - 1) / 2.0
x_keep = np.logical_and(x >= 0, x < resolution)
y_keep = np.logical_and(y >= 0, y < resolution)
keep = np.logical_and(x_keep, y_keep)
x = x[keep].astype(np.int32)
y = y[keep].astype(np.int32)
target_img = np.zeros((resolution, resolution), dtype=np.uint8)
target_img[y, x] = 255
plt.subplot(2, 3, 6)
plt.imshow(target_img)
plt.title("Init Mesh")
if not vis_fg:
mask = proposals[i].gt_masks[j].cpu().numpy()
plt.subplot(1, 4, 2)
plt.imshow(mask)
plt.gca().add_patch(
plt.Rectangle(
(roi[0], roi[1]),
roi[2] - roi[0],
roi[3] - roi[1],
fill=False,
edgecolor="g",
linewidth=3,
)
)
save_file = os.path.join(
output_dir, format(np.random.randint(10000)) + "_" + format(j) + ".png"
)
plt.savefig(save_file)
plt.close(fig)