Skip to content

Commit 4fc6e34

Browse files
lexiyutoulexiyutou
authored
Add render depth in demo and app (#11)
* add render depth n demo.py * change color map * depth map options in app.py and demo_tta.py
1 parent 58dc8b7 commit 4fc6e34

4 files changed

Lines changed: 103 additions & 14 deletions

File tree

‎app.py‎

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -81,8 +81,8 @@
8181
"faster_rcnn_X_101_32x8d_FPN_3x/139173657/model_final_68b088.pkl"
8282
)
8383

84-
# Gradio example row: (image_rel, tta_lr, tta_iters, det_thresh, kp_thresh, side_view, save_mesh)
85-
ExampleRow = Tuple[str, float, int, float, float, bool, bool]
84+
# Gradio example row: (image_rel, tta_lr, tta_iters, det_thresh, kp_thresh, side_view, render_depth, save_mesh)
85+
ExampleRow = Tuple[str, float, int, float, float, bool, bool, bool]
8686

8787

8888
@dataclass(frozen=True)
@@ -98,6 +98,7 @@ class DemoProfile:
9898
max_tta_iters: int
9999
default_save_mesh: bool
100100
default_side_view: bool
101+
default_render_depth: bool
101102
preload_assets: bool
102103
example_rows: Tuple[ExampleRow, ...]
103104
description: str
@@ -124,13 +125,14 @@ def resolve_detectron_device(self) -> str:
124125
max_tta_iters=100,
125126
default_save_mesh=True,
126127
default_side_view=False,
128+
default_render_depth=False,
127129
preload_assets=False,
128130
example_rows=(
129-
("demo_data/000000015956_horse.png", 1e-6, 30, 0.7, 0.1, False, True),
130-
("demo_data/n02412080_12159.png", 1e-6, 30, 0.7, 0.1, False, True),
131-
("demo_data/000000315905_zebra.jpg", 1e-6, 30, 0.7, 0.1, False, True),
132-
("demo_data/beagle.jpg", 1e-6, 30, 0.7, 0.1, False, True),
133-
("demo_data/shepherd_hati.jpg", 1e-6, 30, 0.7, 0.1, False, True),
131+
("demo_data/000000015956_horse.png", 1e-6, 30, 0.7, 0.1, False, False, True),
132+
("demo_data/n02412080_12159.png", 1e-6, 30, 0.7, 0.1, False, False, True),
133+
("demo_data/000000315905_zebra.jpg", 1e-6, 30, 0.7, 0.1, False, False, True),
134+
("demo_data/beagle.jpg", 1e-6, 30, 0.7, 0.1, False, False, True),
135+
("demo_data/shepherd_hati.jpg", 1e-6, 30, 0.7, 0.1, False, False, True),
134136
),
135137
description=(
136138
"**Local demo** — full pipeline on your machine (GPU when available).\n\n"
@@ -153,13 +155,14 @@ def resolve_detectron_device(self) -> str:
153155
max_tta_iters=30,
154156
default_save_mesh=False,
155157
default_side_view=False,
158+
default_render_depth=False,
156159
preload_assets=True,
157160
example_rows=(
158-
("demo_data/000000015956_horse.png", 1e-6, 30, 0.7, 0.1, False, False),
159-
("demo_data/n02412080_12159.png", 1e-6, 30, 0.7, 0.1, False, False),
160-
("demo_data/000000315905_zebra.jpg", 1e-6, 30, 0.7, 0.1, False, False),
161-
("demo_data/beagle.jpg", 1e-6, 30, 0.7, 0.1, False, False),
162-
("demo_data/shepherd_hati.jpg", 1e-6, 30, 0.7, 0.1, False, False),
161+
("demo_data/000000015956_horse.png", 1e-6, 30, 0.7, 0.1, False, False, False),
162+
("demo_data/n02412080_12159.png", 1e-6, 30, 0.7, 0.1, False, False, False),
163+
("demo_data/000000315905_zebra.jpg", 1e-6, 30, 0.7, 0.1, False, False, False),
164+
("demo_data/beagle.jpg", 1e-6, 30, 0.7, 0.1, False, False, False),
165+
("demo_data/shepherd_hati.jpg", 1e-6, 30, 0.7, 0.1, False, False, False),
163166
),
164167
description=(
165168
"**Hugging Face Space (cpu-basic)** — lightweight demo: **CPU-only** PRIMA inference. "
@@ -447,6 +450,7 @@ def _collect_animal_results(
447450
det_thresh: float,
448451
kp_conf_thresh: float,
449452
side_view: bool,
453+
render_depth: bool,
450454
save_mesh: bool,
451455
boxes: Optional[np.ndarray] = None,
452456
progress_callback: Optional[Callable[[str], None]] = None,
@@ -526,6 +530,7 @@ def report(message: str) -> None:
526530
suffix="before_tta",
527531
side_view=side_view,
528532
save_mesh=save_mesh,
533+
render_depth=render_depth,
529534
)
530535

531536
before_png_path = os.path.join(out_folder, f"{img_fn}_{animal_id}_before_tta.png")
@@ -552,6 +557,7 @@ def report(message: str) -> None:
552557
suffix="after_tta",
553558
side_view=side_view,
554559
save_mesh=save_mesh,
560+
render_depth=render_depth,
555561
)
556562

557563
after_png_path = os.path.join(out_folder, f"{img_fn}_{animal_id}_after_tta.png")
@@ -619,6 +625,7 @@ def report(message: str) -> None:
619625
suffix="after_tta",
620626
side_view=side_view,
621627
save_mesh=save_mesh,
628+
render_depth=render_depth,
622629
)
623630

624631
after_png_path = os.path.join(out_folder, f"{img_fn}_{animal_id}_after_tta.png")
@@ -670,6 +677,7 @@ def gradio_inference(
670677
det_thresh: float,
671678
kp_conf_thresh: float,
672679
side_view: bool,
680+
render_depth: bool,
673681
save_mesh: bool,
674682
):
675683
"""Wrapper for Gradio. ``image`` is an RGB numpy array.
@@ -756,6 +764,7 @@ def run_collect(progress_callback: Optional[Callable[[str], None]] = None):
756764
det_thresh=det_thresh,
757765
kp_conf_thresh=kp_conf_thresh,
758766
side_view=side_view,
767+
render_depth=render_depth,
759768
save_mesh=save_mesh,
760769
boxes=boxes,
761770
progress_callback=progress_callback,
@@ -855,6 +864,7 @@ def report_stage(message: str) -> None:
855864
step=0.05,
856865
),
857866
gr.Checkbox(label="Render side view", value=profile.default_side_view),
867+
gr.Checkbox(label="Render depth map", value=profile.default_render_depth),
858868
gr.Checkbox(label="Save meshes (.obj)", value=profile.default_save_mesh),
859869
],
860870
outputs=[

‎demo.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,8 @@ def main():
5959
parser.add_argument('--out_folder', type=str, default='demo_out', help='Output folder to save rendered results')
6060
parser.add_argument('--side_view', dest='side_view', action='store_true', default=False,
6161
help='If set, render side view also')
62+
parser.add_argument('--render_depth', dest='render_depth', action='store_true', default=False,
63+
help='If set, render depth map also')
6264
parser.add_argument('--save_mesh', dest='save_mesh', action='store_true', default=False,
6365
help='If set, save meshes to disk also')
6466
parser.add_argument('--batch_size', type=int, default=1, help='Batch size for inference/fitting')
@@ -162,6 +164,35 @@ def main():
162164
scene_bg_color=(1, 1, 1),
163165
side_view=True)
164166
final_img = np.concatenate([final_img, side_img], axis=1)
167+
168+
if args.render_depth:
169+
depth_img = renderer(out['pred_vertices'][n].detach().cpu().numpy(),
170+
out['pred_cam_t'][n].detach().cpu().numpy(),
171+
white_img,
172+
mesh_base_color=GREEN,
173+
scene_bg_color=(1, 1, 1),
174+
depth=True)
175+
176+
valid_mask = depth_img > 0
177+
if np.sum(valid_mask) == 0:
178+
# no valid depth
179+
depth_norm = np.zeros_like(depth_img)
180+
else:
181+
min_val = np.min(depth_img[valid_mask])
182+
max_val = np.max(depth_img[valid_mask])
183+
if min_val == max_val:
184+
depth_norm = np.zeros_like(depth_img)
185+
else:
186+
depth_norm = (depth_img - min_val) / (max_val - min_val + 1e-8)
187+
depth_norm[~valid_mask] = 0
188+
189+
depth_vis = (depth_norm * 255).astype(np.uint8)
190+
depth_vis = cv2.applyColorMap(depth_vis, cv2.COLORMAP_VIRIDIS)
191+
depth_vis = cv2.cvtColor(depth_vis, cv2.COLOR_BGR2RGB)
192+
depth_vis = depth_vis.astype(np.float32) / 255.0
193+
depth_vis[~valid_mask] = 0
194+
final_img = np.concatenate([final_img, depth_vis], axis=1)
195+
165196

166197
cv2.imwrite(os.path.join(args.out_folder, f'{img_fn}_{animal_id}.png'),
167198
cv2.cvtColor((255 * final_img).astype(np.uint8), cv2.COLOR_RGB2BGR))

‎demo_tta.py‎

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,27 @@ def save_keypoint_vis(patch_rgb: np.ndarray, kpts_xyc: np.ndarray, save_path: st
9292
cv2.imwrite(save_path, vis)
9393

9494

95+
def depth_to_viridis_rgb(depth_img: np.ndarray) -> np.ndarray:
96+
valid_mask = depth_img > 0
97+
if np.sum(valid_mask) == 0:
98+
depth_norm = np.zeros_like(depth_img)
99+
else:
100+
min_val = np.min(depth_img[valid_mask])
101+
max_val = np.max(depth_img[valid_mask])
102+
if min_val == max_val:
103+
depth_norm = np.zeros_like(depth_img)
104+
else:
105+
depth_norm = (depth_img - min_val) / (max_val - min_val + 1e-8)
106+
depth_norm[~valid_mask] = 0
107+
108+
depth_vis = (depth_norm * 255).astype(np.uint8)
109+
depth_vis = cv2.applyColorMap(depth_vis, cv2.COLORMAP_VIRIDIS)
110+
depth_vis = cv2.cvtColor(depth_vis, cv2.COLOR_BGR2RGB)
111+
depth_vis = depth_vis.astype(np.float32) / 255.0
112+
depth_vis[~valid_mask] = 0
113+
return depth_vis
114+
115+
95116
def resolve_sa_weights_path(local_path: str) -> str:
96117
"""Return a local path to the fine-tuned SuperAnimal .pt snapshot.
97118
@@ -162,7 +183,19 @@ def run_superanimal_on_patch(patch_rgb: np.ndarray, args, tmp_dir: str):
162183
return bodyparts[best_idx].astype(np.float32)
163184

164185

165-
def render_and_save(renderer, cam_crop_to_full_fn, out, batch, img_fn, animal_id, out_folder, suffix, side_view, save_mesh):
186+
def render_and_save(
187+
renderer,
188+
cam_crop_to_full_fn,
189+
out,
190+
batch,
191+
img_fn,
192+
animal_id,
193+
out_folder,
194+
suffix,
195+
side_view,
196+
save_mesh,
197+
render_depth=False,
198+
):
166199
pred_cam = out['pred_cam']
167200
box_center = batch['box_center'].float()
168201
box_size = batch['box_size'].float()
@@ -195,6 +228,17 @@ def render_and_save(renderer, cam_crop_to_full_fn, out, batch, img_fn, animal_id
195228
)
196229
final_img = np.concatenate([final_img, side_img], axis=1)
197230

231+
if render_depth:
232+
depth_img = renderer(
233+
out['pred_vertices'][0].detach().cpu().numpy(),
234+
out['pred_cam_t'][0].detach().cpu().numpy(),
235+
white_img,
236+
mesh_base_color=GREEN,
237+
scene_bg_color=(1, 1, 1),
238+
depth=True,
239+
)
240+
final_img = np.concatenate([final_img, depth_to_viridis_rgb(depth_img)], axis=1)
241+
198242
cv2.imwrite(
199243
os.path.join(out_folder, f'{img_fn}_{animal_id}_{suffix}.png'),
200244
cv2.cvtColor((255 * final_img).astype(np.uint8), cv2.COLOR_RGB2BGR),
@@ -252,6 +296,7 @@ def main():
252296
parser.add_argument('--img_folder', type=str, default='demo_data/', help='Folder with input images')
253297
parser.add_argument('--out_folder', type=str, default='demo_out_tta', help='Output folder')
254298
parser.add_argument('--side_view', dest='side_view', action='store_true', default=False, help='Render side view')
299+
parser.add_argument('--render_depth', dest='render_depth', action='store_true', default=False, help='Render depth map')
255300
parser.add_argument('--save_mesh', dest='save_mesh', action='store_true', default=False, help='Save meshes')
256301
parser.add_argument('--file_type', nargs='+', default=['*.jpg', '*.png', '*.jpeg', '*.JPEG'], help='Image globs')
257302
parser.add_argument('--det_thresh', type=float, default=0.7, help='Detectron2 score threshold for animals')
@@ -348,6 +393,7 @@ def main():
348393
suffix='before_tta',
349394
side_view=args.side_view,
350395
save_mesh=args.save_mesh,
396+
render_depth=args.render_depth,
351397
)
352398

353399
patch_rgb = denorm_patch_to_rgb(batch['img'][0])
@@ -392,6 +438,7 @@ def main():
392438
suffix='after_tta',
393439
side_view=args.side_view,
394440
save_mesh=args.save_mesh,
441+
render_depth=args.render_depth,
395442
)
396443

397444

‎demo_tta.sh‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,4 +12,5 @@ python3 demo_tta.py \
1212
--img_folder demo_data/ \
1313
--out_folder demo_out_tta/ \
1414
--tta_lr 1e-6 \
15-
--tta_num_iters 30
15+
--tta_num_iters 30 \
16+
--render_depth

0 commit comments

Comments
 (0)