adds depth, anatomy segmentation viewer

gradio_app.py

@@ -1,6 +1,6 @@
 from functools import partial
 import gradio as gr
-
+import pdb
 from PIL import Image
 import numpy as np
 import gradio as gr
@@ -129,6 +129,59 @@ background_ds = Background2d(
     image_filenames=None,
 )
 
+from dermsynth3d.utils.anatomy import SimpleAnatomy
+color_labels = {
+            0: (0., 0., 0.),  # background
+    1: (174., 199., 232.), # head
+    2: (152., 223., 138.), # torso
+    3: (31., 119., 180.), # hips
+    4: (255., 187., 120.), # legs
+    5: (188., 189., 34.), # feet
+    6: (140., 86., 75.), # arms
+    7: (255., 152., 150.), # hands
+    }
+
+
+def to_simple_anatomy(anatomy):
+    for i in range(16+1):
+        if i in [0,1]:
+            continue
+        if i in [2,3]:
+            anatomy[anatomy==i] = 2
+        if i == 4:
+            anatomy[anatomy==i] = 3
+        if i in [5,6,7,8]:
+            anatomy[anatomy==i] = 4
+        if i in [9,10]:
+            anatomy[anatomy==i] = 5
+        if i in [11,12,13,14]:
+            anatomy[anatomy==i] = 6
+        if i in [15,16]:
+            anatomy[anatomy==i] = 7
+    return anatomy
+
+def convert_anatomy_to_rgb(anatomy):
+    anatomy = to_simple_anatomy(anatomy)
+    anatomy_rgb = np.zeros((anatomy.shape[0], anatomy.shape[1], 3))
+    for k, v in color_labels.items():
+        anatomy_rgb[anatomy == k] = v
+    return anatomy_rgb.astype(np.uint8)
+
+import PIL.Image as pil
+import numpy as np
+import matplotlib as mpl
+import matplotlib.cm as cm
+def convert_depth_to_rgb(depth):
+    mask = depth != 0
+    disp_map = 1 / depth
+    vmax = np.percentile(disp_map[mask], 95)
+    vmin = np.percentile(disp_map[mask], 5)
+    normalizer = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
+    mapper = cm.ScalarMappable(norm=normalizer, cmap='magma')
+    mask = np.repeat(np.expand_dims(mask,-1), 3, -1)
+    colormapped_im = (mapper.to_rgba(disp_map)[:, :, :3] * 255).astype(np.uint8)
+    colormapped_im[~mask] = 255
+    return colormapped_im
 
 def prepare_ds_renderer(
     randomize,
@@ -168,30 +221,30 @@ def prepare_ds_renderer(
         config=None,
         is_blended=is_blend,
     )
-    blended3d = Blended3d(
-        mesh_filename=os.path.join(
-            "./DermSynth3D/data/3dbodytex-1.1-highres/",
-            mesh_name,
-            "model_highres_0_normalized.obj",
-        ),
-        device=DEVICE,
-        dir_blended_textures=dir_blended_textures,
-        dir_anatomy=dir_anatomy,
-        extension=extension ,
-    )
-    normal_texture = load_texture_map(
-        mesh, mesh_name, "No Lesion", 0, device
-    ).maps_padded()
-    if num_lesion > 0:
-        blended_texture_image = load_texture_map(
-            mesh, mesh_name, "Blended Lesion", num_lesion, device
-        ).maps_padded()
-        pasted_texture_image = load_texture_map(
-            mesh, mesh_name, "Pasted Lesion", num_lesion, device
-        ).maps_padded()
-        dilated_texture_image = load_texture_map(
-            mesh, mesh_name, "Dilated Lesion", num_lesion, device
-        ).maps_padded()
+    # blended3d = Blended3d(
+    #     mesh_filename=os.path.join(
+    #         "./DermSynth3D/data/3dbodytex-1.1-highres/",
+    #         mesh_name,
+    #         "model_highres_0_normalized.obj",
+    #     ),
+    #     device=DEVICE,
+    #     dir_blended_textures=dir_blended_textures,
+    #     dir_anatomy=dir_anatomy,
+    #     extension=extension ,
+    # )
+    # normal_texture = load_texture_map(
+    #     mesh, mesh_name, "No Lesion", 0, device
+    # ).maps_padded()
+    # if num_lesion > 0:
+    #     blended_texture_image = load_texture_map(
+    #         mesh, mesh_name, "Blended Lesion", num_lesion, device
+    #     ).maps_padded()
+    #     pasted_texture_image = load_texture_map(
+    #         mesh, mesh_name, "Pasted Lesion", num_lesion, device
+    #     ).maps_padded()
+    #     dilated_texture_image = load_texture_map(
+    #         mesh, mesh_name, "Dilated Lesion", num_lesion, device
+    #     ).maps_padded()
 
     # texture_lesion_mask = blended3d.lesion_texture_mask(astensor=True).to(device)
     # non_skin_texture_mask = blended3d.nonskin_texture_mask(astensor=True).to(device)
@@ -210,16 +263,8 @@ def prepare_ds_renderer(
         mat_sh,
         mat_sc,
     )
-    # mesh_renderer.mesh = mesh
-    # mesh_renderer.cameras = cameras
-    # mesh_renderer.lights = lights
-    # mesh_renderer.materials = materials
-    # mesh_renderer.renderer = renderer
     gr.Info("Successfully prepared renderer.")
-    # render normal images
     gr.Info("Rendering Images...")
-    # if num_views > 1:
-    #     mesh_renderer.mesh = mesh.extend(num_views)
     gr.Info(f"Rendering {num_views} views on {DEVICE}. Please wait...")
     img_count = 0
     view2d = []
@@ -227,10 +272,11 @@ def prepare_ds_renderer(
     anatomy2d = []
     seg2d = []
     view_size = (224, 224)
+    gen2d.view_size = view_size
     while img_count < num_views:
         if randomize:
             gr.Info("Finding suitable parameters...")
-            success = gen2d.randomize_parameters(config=None)
+            success = gen2d.randomize_parameters(config=None, view_size=view_size)
             if not success:
                 gr.Info("Could not find suitable parameters. Trying again.")
                 continue
@@ -238,9 +284,7 @@ def prepare_ds_renderer(
             raster_settings = RasterizationSettings(
                 image_size=view_size[0],
                 blur_radius=0.0,
-                faces_per_pixel=1,
-                # max_faces_per_bin=100,
-                # bin_size=0,
+                faces_per_pixel=10,
                 perspective_correct=True,
             )
             gen2d.mesh_renderer.cameras = cameras
@@ -259,60 +303,19 @@ def prepare_ds_renderer(
             print("***Not enough skin or unable to paste lesion. Skipping.")
             continue
         paste_img = (paste_img * 255).astype(np.uint8)
+        anatomy_view = target[:, :, 3]
         depth_view = target[:, :, 4]
-        depth_img = (depth_view - depth_view.min()) / (
-            depth_view.max() - depth_view.min()
-        )
-        depth_img = (depth_img * 255).astype(np.uint8)
+        depth_img = convert_depth_to_rgb(depth_view)
         view2d.append(paste_img)
         depth2d.append(depth_img)
-        anatomy2d.append(target[:, :, 5])
-        seg2d.append(target[:, :, 3])
+        anatomy_img = convert_anatomy_to_rgb(anatomy_view)
+        anatomy2d.append(anatomy_img)
+        mask = target[:, :, 0]
+        seg2d.append(mask)
         gr.Info(f"Successfully rendered {img_count+1}/{num_views} image+annotations.")
         img_count += 1
     return view2d, depth2d, anatomy2d, seg2d
 
-    # mesh_renderer.compute_fragments()
-    # view2d = mesh_renderer.render_view(asnumpy=True, asRGB=True)
-    # gr.Info("Successfully rendered images.")
-    # gr.Info("Preparing annotations...")
-    # # breakpoint()
-    # pix2face = torch.from_numpy(mesh_renderer.pixels_to_face()).to(
-    #     mesh_renderer.mesh.device
-    # )
-    # pix2vert = torch.stack(
-    #     [a[i] for a, i in zip(mesh_renderer.mesh.faces_padded().squeeze(), pix2face)]
-    # )
-    # pix2vert = pix2vert.detach().cpu().numpy()
-    # anatomy_image = [
-    #     vertices_to_anatomy[pix2vert[i]] * mesh_renderer.body_mask()
-    #     for i in range(num_views)
-    # ]
-    # anatomy_image = np.stack(anatomy_image)
-
-    # anatomy_image = mesh_renderer.anatomy_image(vertices_to_anatomy)
-    # depth_img = mesh_renderer.depth_view(asnumpy=True)
-    # mesh_renderer.set_texture_image(texture_lesion_mask[:, :, np.newaxis])
-    # mask2d = mesh_renderer.render_view(asnumpy=True, asRGB=True)
-    # lesion_mask = mesh_renderer.lesion_mask(mask2d[:, :, 0], lesion_mask_id=None)
-    # # skin mask
-    # mesh_renderer.set_texture_image(non_skin_texture_mask)
-    # nonskin_mask = mesh_renderer.render_view(asnumpy=True, asRGB=True)
-    # skin_mask = mesh_renderer.skin_mask(nonskin_mask[:, :, 0] > 0.5)
-    # segmentation_mask = make_masks(lesion_mask, skin_mask)
-    # gr.Info("Successfully prepared annotations.")
-    # print(view2d.shape, anatomy_image.shape, depth_img.shape, segmentation_mask.shape)
-    # convert anatomy image with labels for each pixel to an image with RGB values
-    # map labels to pixels
-
-    # return (
-    #     view2d,
-    #     anatomy_image,
-    #     depth_img,
-    #     skin_mask,
-    # )  # segmentation_mask
-
-
 # define the list of all the examples
 def get_examples():
     # setup_paths()
@@ -392,7 +395,7 @@ def plotly_mesh(verts, faces, vc, mesh_name):
             )
         ]
     )
-    fig.update_layout(scene_aspectmode="manual", scene_aspectratio=dict(x=1, y=1, z=1))
+    # fig.update_layout(scene_aspectmode="manual", scene_aspectratio=dict(x=1, y=1, z=1))
     fig.update_layout(scene=dict(xaxis=dict(visible=False), yaxis=dict(visible=False)))
     fig.update_layout(scene=dict(zaxis=dict(visible=False)))
     fig.update_layout(scene=dict(camera=dict(up=dict(x=1, y=0, z=1))))
@@ -624,7 +627,6 @@ def process_examples(mesh_name, tex_name, n_lesion):
     mesh_path = get_mesh_path(mesh_name)
     texture_path = get_texture_module(tex_name)(mesh_name, n_lesion)
     mesh_to_view = plotly_mesh(*get_trimesh_attrs(mesh_name, tex_name, n_lesion))
-    # mesh = load_mesh_and_texture(mesh_name, tex_name, n_lesion)
     return mesh_to_view, texture_path, n_lesion
 
 
@@ -640,7 +642,7 @@ def update_plots(mesh_name, texture_name, num_lesion):
         )
         return default_mesh_plot, default_texture, num_lesion
     mesh_path = get_mesh_path(mesh_name)
-    texture_path = get_texture_module(texture_name)(mesh_name, num_lesion)
+    texture_path = Image.open(get_texture_module(texture_name)(mesh_name, num_lesion)).convert("RGB").resize((512, 512))
     mesh_to_view = plotly_mesh(*get_trimesh_attrs(mesh_name, texture_name, num_lesion))
     gr.Info("Successfully updated mesh and texture.")
     return mesh_to_view, texture_path, num_lesion
@@ -852,9 +854,26 @@ def run_demo():
         )
         # rendered views panel
         with gr.Row(variant="panel"):
-            render_block = gr.Gallery(
-                label="Rendered Views", columns=4, height="auto", object_fit="contain"
-            )
+            with gr.Tab("Rendered RGB Views"):
+                render_block = gr.Gallery(
+                    label="Rendered Views", columns=4, height="auto", object_fit="contain"
+                )
+            with gr.Tab("Rendered Depth Views"):
+                depth_block = gr.Gallery(
+                    label="Depth Maps", columns=4, height="auto", object_fit="contain"
+                )
+            with gr.Tab("Rendered Anatomy Views"):
+                anatomy_block = gr.Gallery(
+                    label="Anatomy Labels", columns=4, height="auto", object_fit="contain"
+                )
+            with gr.Tab("Rendered Segmentation Views"):
+                seg_block = gr.Gallery(
+                    label="Segmentation Masks", columns=4, height="auto", object_fit="contain"
+                )
+            #
+            # render_block = gr.Gallery(
+            #     label="Rendered Views", columns=4, height="auto", object_fit="contain"
+            # )
 
         @gr.on(
             triggers=[render_button.click],
@@ -874,7 +893,7 @@ def run_demo():
                 mat_sh,
                 mat_sc,
             ],
-            outputs=[render_block],
+            outputs=[render_block, depth_block, anatomy_block, seg_block],
         )
         def render_views(
             randomize,
@@ -912,7 +931,7 @@ def run_demo():
             #     renderer, mesh, lights, cameras, materials, num_views
             # )
             # return [_ for _ in images.detach().cpu().numpy()]
-            view2d, anatomy, depth, segmentation = prepare_ds_renderer(
+            view2d, depth, anatomy, segmentation = prepare_ds_renderer(
                 randomize,
                 select_mesh,
                 select_texture,
@@ -928,7 +947,7 @@ def run_demo():
                 mat_sh,
                 mat_sc,
             )
-            return view2d
+            return view2d, depth, anatomy, segmentation
 
         # examples panel when the iuser does not want to input
         with gr.Row(variant="panel"):