Fix overlay rendering in fit_single_frame

vchoutas · vchoutas · commit f97a01f98c39 · 2019-08-14T10:06:10.000+02:00
diff --git a/smplifyx/fit_single_frame.py b/smplifyx/fit_single_frame.py
@@ -511,7 +511,7 @@ def fit_single_frame(img,
 
         material = pyrender.MetallicRoughnessMaterial(
             metallicFactor=0.0,
-            alphaMode='BLEND',
+            alphaMode='OPAQUE',
             baseColorFactor=(1.0, 1.0, 0.9, 1.0))
         mesh = pyrender.Mesh.from_trimesh(
             out_mesh,
@@ -523,6 +523,9 @@ def fit_single_frame(img,
 
         camera_center = camera.center.detach().cpu().numpy().squeeze()
         camera_transl = camera.translation.detach().cpu().numpy().squeeze()
+        # Equivalent to 180 degrees around the y-axis. Transforms the fit to
+        # OpenGL compatible coordinate system.
+        camera_transl[0] *= -1.0
 
         camera_pose = np.eye(4)
         camera_pose[:3, 3] = camera_transl
@@ -532,26 +535,21 @@ def fit_single_frame(img,
             cx=camera_center[0], cy=camera_center[1])
         scene.add(camera, pose=camera_pose)
 
-        light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=1)
-
-        light_pose = np.eye(4)
-        light_pose[:3, 3] = [0, -1, 1]
-        scene.add(light, pose=light_pose)
-
-        light_pose[:3, 3] = [0, 1, 1]
-        scene.add(light, pose=light_pose)
-
-        light_pose[:3, 3] = [1, 1, 2]
-        scene.add(light, pose=light_pose)
+        # Get the lights from the viewer
+        light_nodes = monitor.mv.viewer._create_raymond_lights()
+        for node in light_nodes:
+            scene.add_node(node)
 
         r = pyrender.OffscreenRenderer(viewport_width=W,
                                        viewport_height=H,
                                        point_size=1.0)
-        color, _ = r.render(scene)
-
-        color = color.astype(np.float32) / 255
+        color, _ = r.render(scene, flags=pyrender.RenderFlags.RGBA)
+        color = color.astype(np.float32) / 255.0
 
-        output_img = color.copy()
+        valid_mask = (color[:, :, -1] > 0)[:, :, np.newaxis]
+        input_img = img.detach().cpu().numpy()
+        output_img = (color[:, :, :-1] * valid_mask +
+                      (1 - valid_mask) * input_img)
 
         img = pil_img.fromarray((output_img * 255).astype(np.uint8))
         img.save(out_img_fn)
